From 59678ea61097300404ddc40331c57653f8800999 Mon Sep 17 00:00:00 2001
From: Nikita Kryuchkov <nkryuchkov10@gmail.com>
Date: Thu, 20 Jun 2019 18:42:59 +0300
Subject: [PATCH] Implement a test if a client reconnects to a server

---
 pkg/dmsg/client.go      |  6 ++-
 pkg/dmsg/server_test.go | 91 +++++++++++++++++++++++++++--------------
 2 files changed, 65 insertions(+), 32 deletions(-)

diff --git a/pkg/dmsg/client.go b/pkg/dmsg/client.go
index 1dc4ab210f..0486dd8705 100644
--- a/pkg/dmsg/client.go
+++ b/pkg/dmsg/client.go
@@ -19,6 +19,10 @@ import (
 	"github.com/skycoin/skywire/pkg/transport"
 )
 
+const (
+	clientReconnectInterval = 3 * time.Second
+)
+
 var (
 	// ErrNoSrv indicate that remote client does not have DelegatedServers in entry.
 	ErrNoSrv = errors.New("remote has no DelegatedServers")
@@ -440,7 +444,7 @@ func (c *Client) findOrConnectToServer(ctx context.Context, srvPK cipher.PubKey)
 		select {
 		case <-c.done:
 		case <-ctx.Done():
-		case <-time.After(time.Second * 3):
+		case <-time.After(clientReconnectInterval):
 			conn.log.WithField("remoteServer", srvPK).Warn("Reconnecting")
 			if _, err := c.findOrConnectToServer(ctx, srvPK); err != nil {
 				conn.log.WithError(err).WithField("remoteServer", srvPK).Warn("ReconnectionFailed")
diff --git a/pkg/dmsg/server_test.go b/pkg/dmsg/server_test.go
index 9758ca8bfd..1302191f57 100644
--- a/pkg/dmsg/server_test.go
+++ b/pkg/dmsg/server_test.go
@@ -652,49 +652,78 @@ func TestServer_Serve(t *testing.T) {
 	})
 
 	t.Run("Reconnect to server should succeed", func(t *testing.T) {
-		// generate keys for both clients
-		aPK, aSK := cipher.GenerateKeyPair()
-		bPK, bSK := cipher.GenerateKeyPair()
+		t.Run("Same address", func(t *testing.T) {
+			t.Parallel()
+			testReconnect(t, false)
+		})
 
-		assert.Equal(t, 0, s.connCount())
+		t.Run("Random address", func(t *testing.T) {
+			t.Parallel()
+			testReconnect(t, true)
+		})
+	})
+}
 
-		// create remote
-		a := NewClient(aPK, aSK, dc)
-		a.SetLogger(logging.MustGetLogger("A"))
-		err = a.InitiateServerConnections(context.Background(), 1)
-		require.NoError(t, err)
+func testReconnect(t *testing.T, randomAddr bool) {
+	const smallDelay = 100 * time.Millisecond
 
-		assert.Equal(t, 1, s.connCount())
+	sPK, sSK := cipher.GenerateKeyPair()
+	dc := client.NewMock()
 
-		// create initiator
-		b := NewClient(bPK, bSK, dc)
-		b.SetLogger(logging.MustGetLogger("B"))
-		err = b.InitiateServerConnections(context.Background(), 1)
-		require.NoError(t, err)
+	l, err := nettest.NewLocalListener("tcp")
+	require.NoError(t, err)
+
+	s, err := NewServer(sPK, sSK, "", l, dc)
+	require.NoError(t, err)
 
-		time.Sleep(5 * time.Second)
-		assert.Equal(t, 2, s.connCount())
+	serverAddr := s.Addr()
 
-		err := s.Close()
-		assert.NoError(t, err)
+	go s.Serve() // nolint:errcheck
 
-		time.Sleep(5 * time.Second)
+	aPK, aSK := cipher.GenerateKeyPair()
+	bPK, bSK := cipher.GenerateKeyPair()
 
-		assert.Equal(t, 0, s.connCount())
+	assert.Equal(t, 0, s.connCount())
 
-		// s, err = NewServer(sPK, sSK, s.addr, l, dc)
-		// require.NoError(t, err)
+	remote := NewClient(aPK, aSK, dc)
+	remote.SetLogger(logging.MustGetLogger("remote"))
+	err = remote.InitiateServerConnections(context.Background(), 1)
+	require.NoError(t, err)
 
-		go func() {
-			if err := s.Serve(); err != nil {
-				log.Fatal(err)
-			}
-		}()
+	time.Sleep(smallDelay)
+	assert.Equal(t, 1, s.connCount())
 
-		time.Sleep(5 * time.Second)
+	initiator := NewClient(bPK, bSK, dc)
+	initiator.SetLogger(logging.MustGetLogger("initiator"))
+	err = initiator.InitiateServerConnections(context.Background(), 1)
+	require.NoError(t, err)
 
-		assert.Equal(t, 2, s.connCount())
-	})
+	time.Sleep(smallDelay)
+	assert.Equal(t, 2, s.connCount())
+
+	err = s.Close()
+	assert.NoError(t, err)
+
+	assert.Equal(t, 0, s.connCount())
+
+	addr := ""
+	if !randomAddr {
+		addr = serverAddr
+	}
+
+	l, err = net.Listen("tcp", serverAddr)
+	require.NoError(t, err)
+
+	s, err = NewServer(sPK, sSK, addr, l, dc)
+	require.NoError(t, err)
+
+	go s.Serve() // nolint:errcheck
+
+	time.Sleep(clientReconnectInterval + smallDelay)
+	assert.Equal(t, 2, s.connCount())
+
+	err = s.Close()
+	assert.NoError(t, err)
 }
 
 // Given two client instances (a & b) and a server instance (s),