From dcf4fed6a3384f185e9b1ba034342e086b17542e Mon Sep 17 00:00:00 2001
From: Asim Aslam <asim@aslam.me>
Date: Wed, 7 Aug 2019 18:56:21 +0100
Subject: [PATCH] Add a second test for two tunnels

---
 tunnel/tunnel_test.go | 61 ++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 54 insertions(+), 7 deletions(-)

diff --git a/tunnel/tunnel_test.go b/tunnel/tunnel_test.go
index 5dd527c5..8577d89f 100644
--- a/tunnel/tunnel_test.go
+++ b/tunnel/tunnel_test.go
@@ -1,13 +1,14 @@
 package tunnel
 
 import (
+	"sync"
 	"testing"
 
 	"github.com/micro/go-micro/transport"
 )
 
 // testAccept will accept connections on the transport, create a new link and tunnel on top
-func testAccept(t *testing.T, tun Tunnel, wait chan bool) {
+func testAccept(t *testing.T, tun Tunnel, wg *sync.WaitGroup) {
 	// listen on some virtual address
 	tl, err := tun.Listen("test-tunnel")
 	if err != nil {
@@ -26,7 +27,7 @@ func testAccept(t *testing.T, tun Tunnel, wait chan bool) {
 		if err := c.Recv(m); err != nil {
 			t.Fatal(err)
 		}
-		close(wait)
+		wg.Done()
 		return
 	}
 }
@@ -38,7 +39,7 @@ func testSend(t *testing.T, tun Tunnel) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	//defer c.Close()
+	defer c.Close()
 
 	m := transport.Message{
 		Header: map[string]string{
@@ -58,16 +59,62 @@ func TestTunnel(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	//defer tun.Close()
+	defer tun.Close()
 
-	wait := make(chan bool)
+	var wg sync.WaitGroup
 
 	// start accepting connections
-	go testAccept(t, tun, wait)
+	wg.Add(1)
+	go testAccept(t, tun, &wg)
 
 	// send a message
 	testSend(t, tun)
 
 	// wait until message is received
-	<-wait
+	wg.Wait()
+}
+
+func TestTwoTunnel(t *testing.T) {
+	// create a new tunnel client
+	tunA := NewTunnel(
+		Address(":9096"),
+		Nodes(":9097"),
+	)
+
+	// create a new tunnel server
+	tunB := NewTunnel(
+		Address(":9097"),
+	)
+
+	// start tunB
+	err := tunB.Connect()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer tunB.Close()
+
+	// start tunA
+	err = tunA.Connect()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer tunA.Close()
+
+	var wg sync.WaitGroup
+
+	// start accepting connections
+	wg.Add(1)
+	go testAccept(t, tunA, &wg)
+
+	wg.Add(1)
+	go testAccept(t, tunB, &wg)
+
+	// send a message
+	testSend(t, tunA)
+
+	// send a message
+	testSend(t, tunB)
+
+	// wait until done
+	wg.Wait()
 }