Try to fix the test

pkg/router/route_group_test.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"math/rand"
 	"strconv"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -419,7 +420,7 @@ func testMultipleWR(t *testing.T, iterations int, rg1, rg2 io.ReadWriter, msg1,
 }
 
 // TODO (Darkren) uncomment and fix
-/*func TestArbitrarySizeOneMessage(t *testing.T) {
+func TestArbitrarySizeOneMessage(t *testing.T) {
 	// Test fails if message size is above 4059
 	const (
 		value1 = 4058 // dmsg/noise.maxFrameSize - 38
@@ -437,9 +438,9 @@ func testMultipleWR(t *testing.T, iterations int, rg1, rg2 io.ReadWriter, msg1,
 
 	wg.Wait()
 
-	t.Run("Value2", func(t *testing.T) {
+	/*t.Run("Value2", func(t *testing.T) {
 		testArbitrarySizeOneMessage(t, value2)
-	})
+	})*/
 }
 
 func TestArbitrarySizeMultipleMessagesByChunks(t *testing.T) {
@@ -506,41 +507,91 @@ func testArbitrarySizeMultipleMessagesByChunks(t *testing.T, size int) {
 }
 
 func testArbitrarySizeOneMessage(t *testing.T, size int) {
+	keys := snettest.GenKeyPairs(2)
+
+	pk1 := keys[0].PK
+	pk2 := keys[1].PK
+
+	// create test env
+	nEnv := snettest.NewEnv(t, keys, []string{stcp.Type})
+
+	tpDisc := transport.NewDiscoveryMock()
+	tpKeys := snettest.GenKeyPairs(2)
+
+	m1, m2, tp1, tp2, err := transport.CreateTransportPair(tpDisc, tpKeys, nEnv, stcp.Type)
+	require.NoError(t, err)
+	require.NotNil(t, tp1)
+	require.NotNil(t, tp2)
+	require.NotNil(t, tp1.Entry)
+	require.NotNil(t, tp2.Entry)
+
+	rg0 := createRouteGroup()
 	rg1 := createRouteGroup()
-	rg2 := createRouteGroup()
-	m1, m2, teardownEnv := createTransports(t, rg1, rg2, stcp.Type)
+
+	r0RtIDs, err := rg0.rt.ReserveKeys(1)
+	require.NoError(t, err)
+
+	r1RtIDs, err := rg1.rt.ReserveKeys(1)
+	require.NoError(t, err)
+
+	r0FwdRule := routing.ForwardRule(ruleKeepAlive, r0RtIDs[0], r1RtIDs[0], tp1.Entry.ID, pk2, pk1, 0, 0)
+	err = rg0.rt.SaveRule(r0FwdRule)
+	require.NoError(t, err)
+
+	r1FwdRule := routing.ForwardRule(ruleKeepAlive, r1RtIDs[0], r0RtIDs[0], tp2.Entry.ID, pk1, pk2, 0, 0)
+	err = rg1.rt.SaveRule(r1FwdRule)
+	require.NoError(t, err)
+
+	r0FwdRtDesc := r0FwdRule.RouteDescriptor()
+	rg0.desc = r0FwdRtDesc.Invert()
+	rg0.tps = append(rg0.tps, tp1)
+	rg0.fwd = append(rg0.fwd, r0FwdRule)
+
+	r1FwdRtDesc := r1FwdRule.RouteDescriptor()
+	rg1.desc = r1FwdRtDesc.Invert()
+	rg1.tps = append(rg1.tps, tp2)
+	rg1.fwd = append(rg1.fwd, r1FwdRule)
 
 	ctx, cancel := context.WithCancel(context.Background())
 
+	// push close packet from transport to route group
+	go pushPackets(ctx, m2, rg1)
+	go pushPackets(ctx, m1, rg0)
+
+	/*rg1 := createRouteGroup()
+	rg2 := createRouteGroup()
+	m1, m2, teardownEnv := createTransports(t, rg1, rg2, stcp.Type)*/
+
+	//ctx, cancel := context.WithCancel(context.Background())
+
 	defer func() {
 		cancel()
-		teardownEnv()
+		nEnv.Teardown()
 	}()
 
-	go pushPackets(ctx, m1, rg1)
-
-	go pushPackets(ctx, m2, rg2)
-
+	go pushPackets(ctx, m1, rg0)
+	go pushPackets(ctx, m2, rg1)
+
 	msg := []byte(strings.Repeat("A", size))
 
-	_, err := rg1.Write(msg)
+	_, err = rg0.Write(msg)
 	require.NoError(t, err)
 
 	buf := make([]byte, size)
-	n, err := rg2.Read(buf)
+	n, err := rg1.Read(buf)
 	require.NoError(t, err)
 	require.Equal(t, size, n)
 	require.Equal(t, msg, buf)
 
 	buf = make([]byte, size)
-	n, err = rg2.Read(buf)
+	n, err = rg1.Read(buf)
 	require.Equal(t, io.EOF, err)
 	require.Equal(t, 0, n)
 	require.Equal(t, make([]byte, size), buf)
 
+	require.NoError(t, rg0.Close())
 	require.NoError(t, rg1.Close())
-	require.NoError(t, rg2.Close())
-}*/
+}
 
 func TestRouteGroup_LocalAddr(t *testing.T) {
 	rg := createRouteGroup()