- add comments

- avoid double closing messages
- add cleanup routine for relay manager

relay/client/manager_test.go

@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"testing"
+	"time"
 
 	log "github.com/sirupsen/logrus"
 
@@ -47,18 +48,13 @@ func TestForeignConn(t *testing.T) {
 	idAlice := "alice"
 	log.Debugf("connect by alice")
 	clientAlice := NewManager(ctx, addr1, idAlice)
-	err := clientAlice.Serve()
-	if err != nil {
-		t.Fatalf("failed to connect to server: %s", err)
-	}
+	clientAlice.Serve()
 
 	idBob := "bob"
 	log.Debugf("connect by bob")
 	clientBob := NewManager(ctx, addr2, idBob)
-	err = clientBob.Serve()
-	if err != nil {
-		t.Fatalf("failed to connect to server: %s", err)
-	}
+	clientBob.Serve()
+
 	bobsSrvAddr, err := clientBob.RelayAddress()
 	if err != nil {
 		t.Fatalf("failed to get relay address: %s", err)
@@ -137,10 +133,7 @@ func TestForeginConnClose(t *testing.T) {
 	idAlice := "alice"
 	log.Debugf("connect by alice")
 	clientAlice := NewManager(ctx, addr1, idAlice)
-	err := clientAlice.Serve()
-	if err != nil {
-		t.Fatalf("failed to connect to server: %s", err)
-	}
+	clientAlice.Serve()
 
 	conn, err := clientAlice.OpenConn(addr2, "anotherpeer")
 	if err != nil {
@@ -154,3 +147,60 @@ func TestForeginConnClose(t *testing.T) {
 
 	select {}
 }
+
+func TestForeginAutoClose(t *testing.T) {
+	ctx := context.Background()
+
+	addr1 := "localhost:1234"
+	srv1 := server.NewServer()
+	go func() {
+		err := srv1.Listen(addr1)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv1.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	addr2 := "localhost:2234"
+	srv2 := server.NewServer()
+	go func() {
+		err := srv2.Listen(addr2)
+		if err != nil {
+			t.Fatalf("failed to bind server: %s", err)
+		}
+	}()
+
+	defer func() {
+		err := srv2.Close()
+		if err != nil {
+			t.Errorf("failed to close server: %s", err)
+		}
+	}()
+
+	idAlice := "alice"
+	log.Debugf("connect by alice")
+	mgr := NewManager(ctx, addr1, idAlice)
+	relayCleanupInterval = 2 * time.Second
+	mgr.Serve()
+
+	conn, err := mgr.OpenConn(addr2, "anotherpeer")
+	if err != nil {
+		t.Fatalf("failed to bind channel: %s", err)
+	}
+
+	err = conn.Close()
+	if err != nil {
+		t.Fatalf("failed to close connection: %s", err)
+	}
+
+	time.Sleep(relayCleanupInterval + 1*time.Second)
+	if len(mgr.relayClients) != 0 {
+		t.Errorf("expected 0, got %d", len(mgr.relayClients))
+	}
+}