From 0887b6ea33bfa76064ee20476e3b94870ec29859 Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 4 Dec 2023 12:12:06 +0200 Subject: [PATCH 1/3] test(agent): improve TestAgent_Dial tests Refs #11008 --- agent/agent_test.go | 61 +++++++++++++++++++++++++++++++++++++-------- 1 file changed, 51 insertions(+), 10 deletions(-) diff --git a/agent/agent_test.go b/agent/agent_test.go index 31f1448f34018..ba3986a164c66 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "math/rand" @@ -1549,30 +1550,50 @@ func TestAgent_Dial(t *testing.T) { // Setup listener l := c.setup(t) - defer l.Close() + closed := make(chan struct{}) + defer func() { + l.Close() + <-closed + }() + + ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) + defer cancel() + go func() { + var wg sync.WaitGroup + defer func() { + wg.Wait() + close(closed) + }() for { c, err := l.Accept() if err != nil { + if !errors.Is(err, net.ErrClosed) && !errors.Is(err, udp.ErrClosedListener) { + assert.NoError(t, err, "accept connection") + continue + } return } - go testAccept(t, c) + wg.Add(1) + go func() { + defer wg.Done() + testAccept(ctx, t, c) + }() } }() //nolint:dogsled conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0) - require.True(t, conn.AwaitReachable(context.Background())) - conn1, err := conn.DialContext(context.Background(), l.Addr().Network(), l.Addr().String()) + require.True(t, conn.AwaitReachable(ctx)) + conn1, err := conn.DialContext(ctx, l.Addr().Network(), l.Addr().String()) require.NoError(t, err) defer conn1.Close() - conn2, err := conn.DialContext(context.Background(), l.Addr().Network(), l.Addr().String()) + conn2, err := conn.DialContext(ctx, l.Addr().Network(), l.Addr().String()) require.NoError(t, err) defer conn2.Close() - testDial(t, conn2) - testDial(t, conn1) - time.Sleep(150 * time.Millisecond) + testDial(ctx, t, conn2) + testDial(ctx, t, conn1) }) } } @@ -2002,22 +2023,41 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati var dialTestPayload = []byte("dean-was-here123") -func testDial(t *testing.T, c net.Conn) { +func testDial(ctx context.Context, t *testing.T, c net.Conn) { t.Helper() + if deadline, ok := ctx.Deadline(); ok { + err := c.SetDeadline(deadline) + assert.NoError(t, err) + defer func() { + err := c.SetDeadline(time.Time{}) + assert.NoError(t, err) + }() + } + assertWritePayload(t, c, dialTestPayload) assertReadPayload(t, c, dialTestPayload) } -func testAccept(t *testing.T, c net.Conn) { +func testAccept(ctx context.Context, t *testing.T, c net.Conn) { t.Helper() defer c.Close() + if deadline, ok := ctx.Deadline(); ok { + err := c.SetDeadline(deadline) + assert.NoError(t, err) + defer func() { + err := c.SetDeadline(time.Time{}) + assert.NoError(t, err) + }() + } + assertReadPayload(t, c, dialTestPayload) assertWritePayload(t, c, dialTestPayload) } func assertReadPayload(t *testing.T, r io.Reader, payload []byte) { + t.Helper() b := make([]byte, len(payload)+16) n, err := r.Read(b) assert.NoError(t, err, "read payload") @@ -2026,6 +2066,7 @@ func assertReadPayload(t *testing.T, r io.Reader, payload []byte) { } func assertWritePayload(t *testing.T, w io.Writer, payload []byte) { + t.Helper() n, err := w.Write(payload) assert.NoError(t, err, "write payload") assert.Equal(t, len(payload), n, "payload length does not match") From 7f5659f86fc0a89839083de81228797870cd0dba Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 4 Dec 2023 12:20:42 +0200 Subject: [PATCH 2/3] simplify --- agent/agent_test.go | 45 ++++++++++++--------------------------------- 1 file changed, 12 insertions(+), 33 deletions(-) diff --git a/agent/agent_test.go b/agent/agent_test.go index ba3986a164c66..d9b1966bcbde9 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -4,7 +4,6 @@ import ( "bytes" "context" "encoding/json" - "errors" "fmt" "io" "math/rand" @@ -1550,50 +1549,30 @@ func TestAgent_Dial(t *testing.T) { // Setup listener l := c.setup(t) - closed := make(chan struct{}) + done := make(chan struct{}) defer func() { l.Close() - <-closed + <-done }() ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() go func() { - var wg sync.WaitGroup - defer func() { - wg.Wait() - close(closed) - }() - for { - c, err := l.Accept() - if err != nil { - if !errors.Is(err, net.ErrClosed) && !errors.Is(err, udp.ErrClosedListener) { - assert.NoError(t, err, "accept connection") - continue - } - return - } - - wg.Add(1) - go func() { - defer wg.Done() - testAccept(ctx, t, c) - }() - } + defer close(done) + c, err := l.Accept() + assert.NoError(t, err, "accept connection") + defer c.Close() + testAccept(ctx, t, c) }() //nolint:dogsled - conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0) - require.True(t, conn.AwaitReachable(ctx)) - conn1, err := conn.DialContext(ctx, l.Addr().Network(), l.Addr().String()) - require.NoError(t, err) - defer conn1.Close() - conn2, err := conn.DialContext(ctx, l.Addr().Network(), l.Addr().String()) + agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0) + require.True(t, agentConn.AwaitReachable(ctx)) + conn, err := agentConn.DialContext(ctx, l.Addr().Network(), l.Addr().String()) require.NoError(t, err) - defer conn2.Close() - testDial(ctx, t, conn2) - testDial(ctx, t, conn1) + defer conn.Close() + testDial(ctx, t, conn) }) } } From ff6046f77088096fe789f71a4d976ee4bc47030b Mon Sep 17 00:00:00 2001 From: Mathias Fredriksson Date: Mon, 4 Dec 2023 12:24:06 +0200 Subject: [PATCH 3/3] update comment --- agent/agent_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/agent/agent_test.go b/agent/agent_test.go index d9b1966bcbde9..6cef939d47a4c 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -1547,7 +1547,8 @@ func TestAgent_Dial(t *testing.T) { t.Run(c.name, func(t *testing.T) { t.Parallel() - // Setup listener + // The purpose of this test is to ensure that a client can dial a + // listener in the workspace over tailnet. l := c.setup(t) done := make(chan struct{}) defer func() {