fix: code format in master test

shoriwe · Jun 12, 2023 · c09c444 · c09c444
1 parent c42db26
commit c09c444
Showing 1 changed file with 92 additions and 65 deletions.
diff --git a/reverse/master_test.go b/reverse/master_test.go
@@ -2,6 +2,7 @@ package reverse
 
 import (
 	"net"
+	"sync"
 	"testing"
 
 	"github.com/shoriwe/fullproxy/v4/utils/network"
@@ -14,40 +15,40 @@ func TestMaster_init(t *testing.T) {
 		defer data.Close()
 		control := network.ListenAny()
 		defer control.Close()
-		master := network.Dial(control.Addr().String())
-		defer master.Close()
+		masterConn := network.Dial(control.Addr().String())
+		defer masterConn.Close()
 		go func() {
-			s := &Slave{Master: master, Dial: net.Dial}
-			defer s.Close()
-			assert.Nil(tt, s.init())
+			slave := &Slave{Master: masterConn, Dial: net.Dial}
+			defer slave.Close()
+			assert.Nil(tt, slave.init())
 		}()
-		m := &Master{
+		master := &Master{
 			Data:    data,
 			Control: control,
 		}
-		defer m.Close()
-		assert.Nil(tt, m.init())
+		defer master.Close()
+		assert.Nil(tt, master.init())
 	})
 	t.Run("Twice", func(tt *testing.T) {
 		data := network.ListenAny()
 		defer data.Close()
 		control := network.ListenAny()
 		defer control.Close()
+		masterConn := network.Dial(control.Addr().String())
+		defer masterConn.Close()
 		go func() {
-			master := network.Dial(control.Addr().String())
-			defer master.Close()
-			s := &Slave{Master: master, Dial: net.Dial}
-			defer s.Close()
-			assert.Nil(tt, s.init())
-			assert.Nil(tt, s.init())
+			slave := &Slave{Master: masterConn, Dial: net.Dial}
+			defer slave.Close()
+			assert.Nil(tt, slave.init())
+			assert.Nil(tt, slave.init())
 		}()
-		m := &Master{
+		master := &Master{
 			Data:    data,
 			Control: control,
 		}
-		defer m.Close()
-		assert.Nil(tt, m.init())
-		assert.Nil(tt, m.init())
+		defer master.Close()
+		assert.Nil(tt, master.init())
+		assert.Nil(tt, master.init())
 	})
 }
 
@@ -56,42 +57,59 @@ func TestMaster_Addr(t *testing.T) {
 	defer data.Close()
 	control := network.ListenAny()
 	defer control.Close()
-	m := &Master{
+	master := &Master{
 		Data:    data,
 		Control: control,
 	}
-	assert.NotNil(t, m.Addr())
+	defer master.Close()
+	assert.NotNil(t, master.Addr())
 }
 func TestMaster_Accept(t *testing.T) {
 	t.Run("Succeed", func(tt *testing.T) {
 		data := network.ListenAny()
 		defer data.Close()
 		control := network.ListenAny()
 		defer control.Close()
-		master := network.Dial(control.Addr().String())
-		defer master.Close()
-		doneChan := make(chan struct{}, 2)
-		defer close(doneChan)
-		go func() {
-			s := &Slave{Master: master, Dial: net.Dial}
-			defer s.Close()
-			go s.Serve()
-			<-doneChan
-		}()
-		m := &Master{
+		masterConn := network.Dial(control.Addr().String())
+		defer masterConn.Close()
+		// Slave
+		slave := &Slave{Master: masterConn, Dial: net.Dial}
+		defer slave.Close()
+		go slave.Serve()
+		// Master
+		master := &Master{
 			Data:    data,
 			Control: control,
 		}
-		defer m.Close()
+		defer master.Close()
+		// Producer
+		testMsg := []byte("MSG")
+		var wg sync.WaitGroup
+		defer wg.Wait()
+		wg.Add(1)
 		go func() {
-			aConn, aErr := m.Accept()
-			assert.Nil(tt, aErr)
-			defer aConn.Close()
-			<-doneChan
+			defer wg.Done()
+			conn, err := master.Accept()
+			assert.Nil(tt, err)
+			defer conn.Close()
+			// Write
+			_, err = conn.Write(testMsg)
+			assert.Nil(tt, err)
+			// Read
+			buffer := make([]byte, len(testMsg))
+			_, err = conn.Read(buffer)
+			assert.Nil(tt, err)
 		}()
-		aConn := network.Dial(data.Addr().String())
-		defer aConn.Close()
-		doneChan <- struct{}{}
+		// Consumer
+		conn := network.Dial(data.Addr().String())
+		defer conn.Close()
+		// Read
+		buffer := make([]byte, len(testMsg))
+		_, err := conn.Read(buffer)
+		assert.Nil(tt, err)
+		// Write
+		_, err = conn.Write(testMsg)
+		assert.Nil(tt, err)
 	})
 }
 
@@ -101,41 +119,50 @@ func TestMaster_SlaveDial(t *testing.T) {
 		defer data.Close()
 		control := network.ListenAny()
 		defer control.Close()
-		master := network.Dial(control.Addr().String())
-		defer master.Close()
+		masterConn := network.Dial(control.Addr().String())
+		defer masterConn.Close()
 		service := network.ListenAny()
 		defer service.Close()
-		doneChan := make(chan struct{}, 2)
-		defer close(doneChan)
+		// Slave
+		slave := &Slave{Master: masterConn, Dial: net.Dial}
+		defer slave.Close()
+		go slave.Serve()
+		// Master
+		master := &Master{
+			Data:    data,
+			Control: control,
+		}
+		defer master.Close()
+		// Producer
 		testMessage := []byte("TEST")
+		var wg sync.WaitGroup
+		defer wg.Wait()
+		wg.Add(1)
 		go func() {
-			c, err := service.Accept()
+			defer wg.Done()
+			conn, err := service.Accept()
 			assert.Nil(tt, err)
-			defer c.Close()
-			_, err = c.Write(testMessage)
+			defer conn.Close()
+			// Write
+			_, err = conn.Write(testMessage)
+			assert.Nil(tt, err)
+			// Read
+			buffer := make([]byte, len(testMessage))
+			_, err = conn.Read(buffer)
 			assert.Nil(tt, err)
-			<-doneChan
-		}()
-		go func() {
-			s := &Slave{Master: master, Dial: net.Dial}
-			defer s.Close()
-			go s.Serve()
-			<-doneChan
 		}()
-		m := &Master{
-			Data:    data,
-			Control: control,
-		}
-		defer m.Close()
-		serviceConn, dialErr := m.SlaveDial("tcp", service.Addr().String())
-		assert.Nil(tt, dialErr)
-		defer serviceConn.Close()
+		// Consumer
+		conn, err := master.SlaveDial("tcp", service.Addr().String())
+		assert.Nil(tt, err)
+		defer conn.Close()
+		// Read
 		buffer := make([]byte, len(testMessage))
-		_, rErr := serviceConn.Read(buffer)
-		assert.Nil(tt, rErr)
+		_, err = conn.Read(buffer)
+		assert.Nil(tt, err)
 		assert.Equal(tt, testMessage, buffer)
-		doneChan <- struct{}{}
-		doneChan <- struct{}{}
+		// Write
+		_, err = conn.Write(testMessage)
+		assert.Nil(tt, err)
 	})
 	t.Run("Not listening", func(tt *testing.T) {
 		data := network.ListenAny()