diff --git a/pkg/router/route_group_test.go b/pkg/router/route_group_test.go index 7b71d7d04..ba77444a8 100644 --- a/pkg/router/route_group_test.go +++ b/pkg/router/route_group_test.go @@ -60,8 +60,8 @@ func TestRouteGroup_Read(t *testing.T) { rg1 = createRouteGroup() rg2 := createRouteGroup() - _, _, closeFunc := createTransports(t, rg1, rg2) - defer closeFunc() + _, _, teardown := createTransports(t, rg1, rg2) + defer teardown() go func() { rg1.readCh <- msg1 @@ -93,8 +93,8 @@ func TestRouteGroup_Write(t *testing.T) { rg1 = createRouteGroup() rg2 := createRouteGroup() - m1, m2, closeFunc := createTransports(t, rg1, rg2) - defer closeFunc() + m1, m2, teardown := createTransports(t, rg1, rg2) + defer teardown() _, err = rg1.Write(msg1) require.NoError(t, err) @@ -148,13 +148,40 @@ func TestRouteGroup_SetDeadline(t *testing.T) { func TestRouteGroup_TestConn(t *testing.T) { mp := func() (c1, c2 net.Conn, stop func(), err error) { - c1 = createRouteGroup() - c2 = createRouteGroup() + rg1 := createRouteGroup() + rg2 := createRouteGroup() + + c1, c2 = rg1, rg2 + + m1, m2, teardownEnv := createTransports(t, rg1, rg2) + ctx, cancel := context.WithCancel(context.Background()) + + pushPackets := func(from *transport.Manager, to *RouteGroup) { + for { + select { + case <-ctx.Done(): + return + default: + packet, err := from.ReadPacket() + assert.NoError(t, err) + select { + case <-ctx.Done(): + return + case to.readCh <- packet.Payload(): + } + } + } + } + + go pushPackets(m1, rg1) + go pushPackets(m2, rg2) + stop = func() { + cancel() + teardownEnv() assert.NoError(t, c1.Close()) assert.NoError(t, c2.Close()) } - err = nil return } nettest.TestConn(t, mp) @@ -173,7 +200,7 @@ func createRouteGroup() *RouteGroup { return rg } -func createTransports(t *testing.T, rg1, rg2 *RouteGroup) (m1, m2 *transport.Manager, closeFunc func()) { +func createTransports(t *testing.T, rg1, rg2 *RouteGroup) (m1, m2 *transport.Manager, teardown func()) { tpDisc := transport.NewDiscoveryMock() keys := snettest.GenKeyPairs(2)