From d9d92887dce6e42c88b14957a5527f2854e95e3c Mon Sep 17 00:00:00 2001
From: Nikita Kryuchkov <nkryuchkov10@gmail.com>
Date: Thu, 12 Dec 2019 20:59:00 +0300
Subject: [PATCH] Improve Router and RouteGroup tests

---
 pkg/router/route_group.go      |  6 +++++-
 pkg/router/route_group_test.go | 33 +++++++++++++++++++++++++++++++--
 pkg/router/router.go           |  1 -
 pkg/router/router_test.go      | 17 ++++++++++++++---
 4 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/pkg/router/route_group.go b/pkg/router/route_group.go
index 7b4c053d57..b97a2b4b6d 100644
--- a/pkg/router/route_group.go
+++ b/pkg/router/route_group.go
@@ -170,6 +170,10 @@ func (r *RouteGroup) Write(p []byte) (n int, err error) {
 		return 0, timeoutError{}
 	}
 
+	if len(p) == 0 {
+		return 0, nil
+	}
+
 	r.mu.Lock()
 	defer r.mu.Unlock()
 
@@ -252,7 +256,7 @@ func (r *RouteGroup) Close() error {
 
 	r.once.Do(func() {
 		close(r.done)
-		// close(r.readCh) // TODO: uncomment
+		// close(r.readCh) // TODO(nkryuchkov): uncomment
 	})
 
 	return nil
diff --git a/pkg/router/route_group_test.go b/pkg/router/route_group_test.go
index 4c057992d1..139c95f1a0 100644
--- a/pkg/router/route_group_test.go
+++ b/pkg/router/route_group_test.go
@@ -51,7 +51,11 @@ func TestRouteGroup_Read(t *testing.T) {
 	rg1.readCh <- msg1
 	rg2.readCh <- msg2
 
-	n, err := rg1.Read(buf1)
+	n, err := rg1.Read([]byte{})
+	require.Equal(t, 0, n)
+	require.NoError(t, err)
+
+	n, err = rg1.Read(buf1)
 	require.NoError(t, err)
 	require.Equal(t, msg1, buf1)
 	require.Equal(t, len(msg1), n)
@@ -60,6 +64,9 @@ func TestRouteGroup_Read(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, msg2, buf2)
 	require.Equal(t, len(msg2), n)
+
+	require.NoError(t, rg1.Close())
+	require.NoError(t, rg2.Close())
 }
 
 func TestRouteGroup_Write(t *testing.T) {
@@ -79,6 +86,14 @@ func TestRouteGroup_Write(t *testing.T) {
 	m1, m2, teardown := createTransports(t, rg1, rg2, stcp.Type)
 	defer teardown()
 
+	n, err := rg1.Write([]byte{})
+	require.Equal(t, 0, n)
+	require.NoError(t, err)
+
+	n, err = rg2.Write([]byte{})
+	require.Equal(t, 0, n)
+	require.NoError(t, err)
+
 	_, err = rg1.Write(msg1)
 	require.NoError(t, err)
 
@@ -92,6 +107,9 @@ func TestRouteGroup_Write(t *testing.T) {
 	recv, err = m2.ReadPacket()
 	require.NoError(t, err)
 	require.Equal(t, msg1, recv.Payload())
+
+	require.NoError(t, rg1.Close())
+	require.NoError(t, rg2.Close())
 }
 
 func TestRouteGroup_ReadWrite(t *testing.T) {
@@ -121,6 +139,9 @@ func testReadWrite(t *testing.T, iterations int) {
 	assert.NoError(t, rg2.Close())
 
 	teardownEnv()
+
+	require.NoError(t, rg1.Close())
+	require.NoError(t, rg2.Close())
 }
 
 func testRouteGroupReadWrite(t *testing.T, iterations int, rg1, rg2 io.ReadWriter) {
@@ -378,19 +399,25 @@ func testArbitrarySizeOneMessage(t *testing.T, size int) {
 	require.Equal(t, io.EOF, err)
 	require.Equal(t, 0, n)
 	require.Equal(t, make([]byte, size), buf)
+
+	require.NoError(t, rg1.Close())
+	require.NoError(t, rg2.Close())
 }
 
 func TestRouteGroup_LocalAddr(t *testing.T) {
 	rg := createRouteGroup()
 	require.Equal(t, rg.desc.Dst(), rg.LocalAddr())
+
+	require.NoError(t, rg.Close())
 }
 
 func TestRouteGroup_RemoteAddr(t *testing.T) {
 	rg := createRouteGroup()
 	require.Equal(t, rg.desc.Src(), rg.RemoteAddr())
+
+	require.NoError(t, rg.Close())
 }
 
-// TODO: fix hangs
 func TestRouteGroup_TestConn(t *testing.T) {
 	mp := func() (c1, c2 net.Conn, stop func(), err error) {
 		rg1 := createRouteGroup()
@@ -408,6 +435,8 @@ func TestRouteGroup_TestConn(t *testing.T) {
 		stop = func() {
 			cancel()
 			teardownEnv()
+			require.NoError(t, rg1.Close())
+			require.NoError(t, rg2.Close())
 		}
 
 		return
diff --git a/pkg/router/router.go b/pkg/router/router.go
index 23836507b7..fb2e5f4c8c 100644
--- a/pkg/router/router.go
+++ b/pkg/router/router.go
@@ -219,7 +219,6 @@ func (r *router) AcceptRoutes(ctx context.Context) (*RouteGroup, error) {
 	case <-ctx.Done():
 		return nil, ctx.Err()
 	case rules, ok = <-r.accept:
-		break
 	}
 	if !ok {
 		return nil, &net.OpError{Op: "accept", Net: "skynet", Source: nil, Err: errors.New("use of closed network connection")}
diff --git a/pkg/router/router_test.go b/pkg/router/router_test.go
index 8ce1bd5e22..2df277ea22 100644
--- a/pkg/router/router_test.go
+++ b/pkg/router/router_test.go
@@ -88,9 +88,11 @@ func Test_router_AcceptRoutes(t *testing.T) {
 	require.Len(t, rg.tps, 1)
 
 	allRules := rg.rt.AllRules()
-	assert.Len(t, allRules, 2)
-	assert.Contains(t, allRules, fwdRule)
-	assert.Contains(t, allRules, cnsmRule)
+	require.Len(t, allRules, 2)
+	require.Contains(t, allRules, fwdRule)
+	require.Contains(t, allRules, cnsmRule)
+
+	require.NoError(t, r0.Close())
 }
 
 // Ensure that received packets are handled properly in `(*Router).Serve()`.
@@ -122,6 +124,11 @@ func TestRouter_Serve(t *testing.T) {
 
 	// go r1.Serve(context.TODO())
 
+	defer func() {
+		require.NoError(t, r0.Close())
+		require.NoError(t, r1.Close())
+	}()
+
 	// Create dmsg transport between two `snet.Network` entities.
 	tp1, err := rEnv.TpMngrs[1].SaveTransport(context.TODO(), keys[0].PK, dmsg.Type)
 	require.NoError(t, err)
@@ -324,6 +331,10 @@ func TestRouter_Rules(t *testing.T) {
 	r, ok := rIfc.(*router)
 	require.True(t, ok)
 
+	defer func() {
+		require.NoError(t, r.Close())
+	}()
+
 	r.rt = rt
 
 	// CLOSURE: Delete all routing rules.