remove dep on assert lib

testing/scenarios/command_test.go

@@ -3,7 +3,6 @@ package scenarios
 import (
 	"context"
 	"fmt"
-	"io"
 	"testing"
 	"time"
 
@@ -29,12 +28,9 @@ import (
 	"v2ray.com/core/proxy/vmess/inbound"
 	"v2ray.com/core/proxy/vmess/outbound"
 	"v2ray.com/core/testing/servers/tcp"
-	. "v2ray.com/ext/assert"
 )
 
 func TestCommanderRemoveHandler(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -99,30 +95,10 @@ func TestCommanderRemoveHandler(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(clientConfig)
 	common.Must(err)
-
 	defer CloseAllServers(servers)
 
-	{
-		conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
-			IP:   []byte{127, 0, 0, 1},
-			Port: int(clientPort),
-		})
-		if err != nil {
-			t.Fatal(err)
-		}
-		defer conn.Close() // nolint: errcheck
-
-		payload := "commander request."
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		if r := cmp.Diff(response[:nBytes], xor([]byte(payload))); r != "" {
-			t.Fatal(r)
-		}
+	if err := testTCPConn(clientPort, 1024, time.Second*5)(); err != nil {
+		t.Fatal(err)
 	}
 
 	cmdConn, err := grpc.Dial(fmt.Sprintf("127.0.0.1:%d", cmdPort), grpc.WithInsecure(), grpc.WithBlock())
@@ -134,20 +110,22 @@ func TestCommanderRemoveHandler(t *testing.T) {
 		Tag: "d",
 	})
 	common.Must(err)
-	assert(resp, IsNotNil)
+	if resp == nil {
+		t.Error("unexpected nil response")
+	}
 
 	{
 		_, err := net.DialTCP("tcp", nil, &net.TCPAddr{
 			IP:   []byte{127, 0, 0, 1},
 			Port: int(clientPort),
 		})
-		assert(err, IsNotNil)
+		if err == nil {
+			t.Error("unexpected nil error")
+		}
 	}
 }
 
 func TestCommanderAddRemoveUser(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -284,24 +262,10 @@ func TestCommanderAddRemoveUser(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
 	common.Must(err)
+	CloseAllServers(servers)
 
-	{
-		conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
-			IP:   []byte{127, 0, 0, 1},
-			Port: int(clientPort),
-		})
-		common.Must(err)
-
-		payload := "commander request."
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		assert(nBytes, Equals, 0)
-		assert(err, Equals, io.EOF)
-		assert(conn.Close(), IsNil)
+	if err := testTCPConn(clientPort, 1024, time.Second*5)(); err != nil {
+		t.Fatal(err)
 	}
 
 	cmdConn, err := grpc.Dial(fmt.Sprintf("127.0.0.1:%d", cmdPort), grpc.WithInsecure(), grpc.WithBlock())
@@ -323,40 +287,25 @@ func TestCommanderAddRemoveUser(t *testing.T) {
 			}),
 	})
 	common.Must(err)
-	assert(resp, IsNotNil)
+	if resp == nil {
+		t.Fatal("nil response")
+	}
 
-	{
-		conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
-			IP:   []byte{127, 0, 0, 1},
-			Port: int(clientPort),
-		})
-		common.Must(err)
-
-		payload := "commander request."
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-		assert(conn.Close(), IsNil)
+	if err := testTCPConn(clientPort, 1024, time.Second*5)(); err != nil {
+		t.Fatal(err)
 	}
 
 	resp, err = hsClient.AlterInbound(context.Background(), &command.AlterInboundRequest{
 		Tag:       "v",
 		Operation: serial.ToTypedMessage(&command.RemoveUserOperation{Email: "test@v2ray.com"}),
 	})
-	assert(resp, IsNotNil)
 	common.Must(err)
-
-	CloseAllServers(servers)
+	if resp == nil {
+		t.Fatal("nil response")
+	}
 }
 
 func TestCommanderStats(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -515,20 +464,30 @@ func TestCommanderStats(t *testing.T) {
 		Reset_: true,
 	})
 	common.Must(err)
-	assert(sresp.Stat.Name, Equals, name)
-	assert(sresp.Stat.Value, Equals, int64(10240*1024))
+	if r := cmp.Diff(sresp.Stat, &statscmd.Stat{
+		Name:  name,
+		Value: 10240 * 1024,
+	}); r != "" {
+		t.Error(r)
+	}
 
 	sresp, err = sClient.GetStats(context.Background(), &statscmd.GetStatsRequest{
 		Name: name,
 	})
 	common.Must(err)
-	assert(sresp.Stat.Name, Equals, name)
-	assert(sresp.Stat.Value, Equals, int64(0))
+	if r := cmp.Diff(sresp.Stat, &statscmd.Stat{
+		Name:  name,
+		Value: 0,
+	}); r != "" {
+		t.Error(r)
+	}
 
 	sresp, err = sClient.GetStats(context.Background(), &statscmd.GetStatsRequest{
 		Name:   "inbound>>>vmess>>>traffic>>>uplink",
 		Reset_: true,
 	})
 	common.Must(err)
-	assert(sresp.Stat.Value, GreaterThan, int64(10240*1024))
+	if sresp.Stat.Value <= 10240*1024 {
+		t.Error("value < 10240*1024: ", sresp.Stat.Value)
+	}
 }

testing/scenarios/common.go

@@ -160,6 +160,27 @@ func testTCPConn(port net.Port, payloadSize int, timeout time.Duration) func() e
 		}
 		defer conn.Close()
 
+		return testTCPConn2(conn, payloadSize, timeout)()
+	}
+}
+
+func testUDPConn(port net.Port, payloadSize int, timeout time.Duration) func() error {
+	return func() error {
+		conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
+			IP:   []byte{127, 0, 0, 1},
+			Port: int(port),
+		})
+		if err != nil {
+			return err
+		}
+		defer conn.Close()
+
+		return testTCPConn2(conn, payloadSize, timeout)()
+	}
+}
+
+func testTCPConn2(conn net.Conn, payloadSize int, timeout time.Duration) func() error {
+	return func() error {
 		payload := make([]byte, payloadSize)
 		common.Must2(rand.Read(payload))
 

testing/scenarios/dns_test.go

@@ -3,6 +3,7 @@ package scenarios
 import (
 	"fmt"
 	"testing"
+	"time"
 
 	xproxy "golang.org/x/net/proxy"
 	"v2ray.com/core"
@@ -16,12 +17,9 @@ import (
 	"v2ray.com/core/proxy/freedom"
 	"v2ray.com/core/proxy/socks"
 	"v2ray.com/core/testing/servers/tcp"
-	. "v2ray.com/ext/assert"
 )
 
 func TestResolveIP(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -85,24 +83,17 @@ func TestResolveIP(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
 	{
 		noAuthDialer, err := xproxy.SOCKS5("tcp", net.TCPDestination(net.LocalHostIP, serverPort).NetAddr(), nil, xproxy.Direct)
 		common.Must(err)
 		conn, err := noAuthDialer.Dial("tcp", fmt.Sprintf("google.com:%d", dest.Port))
 		common.Must(err)
+		defer conn.Close()
 
-		payload := "test payload"
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-		assert(conn.Close(), IsNil)
+		if err := testTCPConn2(conn, 1024, time.Second*5)(); err != nil {
+			t.Error(err)
+		}
 	}
-
-	CloseAllServers(servers)
 }

testing/scenarios/dokodemo_test.go

@@ -4,6 +4,8 @@ import (
 	"testing"
 	"time"
 
+	"golang.org/x/sync/errgroup"
+
 	"v2ray.com/core"
 	"v2ray.com/core/app/log"
 	"v2ray.com/core/app/proxyman"
@@ -20,7 +22,6 @@ import (
 	"v2ray.com/core/proxy/vmess/outbound"
 	"v2ray.com/core/testing/servers/tcp"
 	"v2ray.com/core/testing/servers/udp"
-	. "v2ray.com/ext/assert"
 )
 
 func TestDokodemoTCP(t *testing.T) {
@@ -121,8 +122,6 @@ func TestDokodemoTCP(t *testing.T) {
 }
 
 func TestDokodemoUDP(t *testing.T) {
-	assert := With(t)
-
 	udpServer := udp.Server{
 		MsgProcessor: xor,
 	}
@@ -198,25 +197,13 @@ func TestDokodemoUDP(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
+	var errg errgroup.Group
 	for port := clientPort; port <= clientPort+clientPortRange; port++ {
-		conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
-			IP:   []byte{127, 0, 0, 1},
-			Port: int(port),
-		})
-		common.Must(err)
-
-		payload := "dokodemo request."
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-		assert(conn.Close(), IsNil)
+		errg.Go(testUDPConn(net.Port(port), 1024, time.Second*5))
+	}
+	if err := errg.Wait(); err != nil {
+		t.Error(err)
 	}
-
-	CloseAllServers(servers)
 }

testing/scenarios/feature_test.go

@@ -33,12 +33,9 @@ import (
 	"v2ray.com/core/testing/servers/tcp"
 	"v2ray.com/core/testing/servers/udp"
 	"v2ray.com/core/transport/internet"
-	. "v2ray.com/ext/assert"
 )
 
 func TestPassiveConnection(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 		SendFirst:    []byte("send first"),
@@ -73,6 +70,7 @@ func TestPassiveConnection(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
 	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
 		IP:   []byte{127, 0, 0, 1},
@@ -84,27 +82,14 @@ func TestPassiveConnection(t *testing.T) {
 		response := make([]byte, 1024)
 		nBytes, err := conn.Read(response)
 		common.Must(err)
-		assert(string(response[:nBytes]), Equals, "send first")
+		if string(response[:nBytes]) != "send first" {
+			t.Error("unexpected first response: ", string(response[:nBytes]))
+		}
 	}
 
-	payload := "dokodemo request."
-	{
-
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
+	if err := testTCPConn2(conn, 1024, time.Second*5)(); err != nil {
+		t.Error(err)
 	}
-
-	{
-		response := make([]byte, 1024)
-		nBytes, err := conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-	}
-
-	assert(conn.Close(), IsNil)
-
-	CloseAllServers(servers)
 }
 
 func TestProxy(t *testing.T) {
@@ -240,8 +225,6 @@ func TestProxy(t *testing.T) {
 }
 
 func TestProxyOverKCP(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -377,30 +360,14 @@ func TestProxyOverKCP(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig, proxyConfig, clientConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
-	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
-		IP:   []byte{127, 0, 0, 1},
-		Port: int(clientPort),
-	})
-	common.Must(err)
-
-	payload := "dokodemo request."
-	nBytes, err := conn.Write([]byte(payload))
-	common.Must(err)
-	assert(nBytes, Equals, len(payload))
-
-	response := make([]byte, 1024)
-	nBytes, err = conn.Read(response)
-	common.Must(err)
-	assert(response[:nBytes], Equals, xor([]byte(payload)))
-	assert(conn.Close(), IsNil)
-
-	CloseAllServers(servers)
+	if err := testTCPConn(clientPort, 1024, time.Second*5)(); err != nil {
+		t.Error(err)
+	}
 }
 
 func TestBlackhole(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -472,35 +439,14 @@ func TestBlackhole(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
-	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
-		IP:   []byte{127, 0, 0, 1},
-		Port: int(serverPort2),
-	})
-	common.Must(err)
-
-	payload := "dokodemo request."
-	{
-
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
+	if err := testTCPConn(serverPort2, 1024, time.Second*5)(); err != nil {
+		t.Error(err)
 	}
-
-	{
-		response := make([]byte, 1024)
-		_, err := conn.Read(response)
-		assert(err, IsNotNil)
-	}
-
-	assert(conn.Close(), IsNil)
-
-	CloseAllServers(servers)
 }
 
 func TestForward(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -542,31 +488,22 @@ func TestForward(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
 	{
 		noAuthDialer, err := xproxy.SOCKS5("tcp", net.TCPDestination(net.LocalHostIP, serverPort).NetAddr(), nil, xproxy.Direct)
 		common.Must(err)
 		conn, err := noAuthDialer.Dial("tcp", "google.com:80")
 		common.Must(err)
+		defer conn.Close()
 
-		payload := "test payload"
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-		assert(conn.Close(), IsNil)
+		if err := testTCPConn2(conn, 1024, time.Second*5)(); err != nil {
+			t.Error(err)
+		}
 	}
-
-	CloseAllServers(servers)
 }
 
 func TestUDPConnection(t *testing.T) {
-	assert := With(t)
-
 	udpServer := udp.Server{
 		MsgProcessor: xor,
 	}
@@ -600,56 +537,20 @@ func TestUDPConnection(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(clientConfig)
 	common.Must(err)
+	CloseAllServers(servers)
 
-	{
-		conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
-			IP:   []byte{127, 0, 0, 1},
-			Port: int(clientPort),
-		})
-		common.Must(err)
-
-		payload := "dokodemo request."
-		for i := 0; i < 5; i++ {
-			nBytes, err := conn.Write([]byte(payload))
-			common.Must(err)
-			assert(nBytes, Equals, len(payload))
-
-			response := make([]byte, 1024)
-			nBytes, err = conn.Read(response)
-			common.Must(err)
-			assert(response[:nBytes], Equals, xor([]byte(payload)))
-		}
-
-		assert(conn.Close(), IsNil)
+	if err := testUDPConn(clientPort, 1024, time.Second*5)(); err != nil {
+		t.Error(err)
 	}
 
 	time.Sleep(20 * time.Second)
 
-	{
-		conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
-			IP:   []byte{127, 0, 0, 1},
-			Port: int(clientPort),
-		})
-		common.Must(err)
-
-		payload := "dokodemo request."
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-		assert(conn.Close(), IsNil)
+	if err := testUDPConn(clientPort, 1024, time.Second*5)(); err != nil {
+		t.Error(err)
 	}
-
-	CloseAllServers(servers)
 }
 
 func TestDomainSniffing(t *testing.T) {
-	assert := With(t)
-
 	sniffingPort := tcp.PickPort()
 	httpPort := tcp.PickPort()
 	serverConfig := &core.Config{
@@ -722,6 +623,7 @@ func TestDomainSniffing(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
 	{
 		transport := &http.Transport{
@@ -736,17 +638,14 @@ func TestDomainSniffing(t *testing.T) {
 
 		resp, err := client.Get("https://www.github.com/")
 		common.Must(err)
-		assert(resp.StatusCode, Equals, 200)
-
-		assert(resp.Write(ioutil.Discard), IsNil)
+		if resp.StatusCode != 200 {
+			t.Error("unexpected status code: ", resp.StatusCode)
+		}
+		common.Must(resp.Write(ioutil.Discard))
 	}
-
-	CloseAllServers(servers)
 }
 
 func TestDialV2Ray(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -822,23 +721,16 @@ func TestDialV2Ray(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
 	client, err := core.New(clientConfig)
 	common.Must(err)
 
 	conn, err := core.Dial(context.Background(), client, dest)
 	common.Must(err)
+	defer conn.Close()
 
-	payload := "commander request."
-	nBytes, err := conn.Write([]byte(payload))
-	common.Must(err)
-	assert(nBytes, Equals, len(payload))
-
-	response := make([]byte, 1024)
-	nBytes, err = conn.Read(response)
-	common.Must(err)
-	assert(response[:nBytes], Equals, xor([]byte(payload)))
-	assert(conn.Close(), IsNil)
-
-	CloseAllServers(servers)
+	if err := testTCPConn2(conn, 1024, time.Second*5)(); err != nil {
+		t.Error(err)
+	}
 }

testing/scenarios/shadowsocks_test.go

@@ -2,7 +2,6 @@ package scenarios
 
 import (
 	"crypto/rand"
-	"sync"
 	"testing"
 	"time"
 
@@ -23,7 +22,6 @@ import (
 	"v2ray.com/core/proxy/shadowsocks"
 	"v2ray.com/core/testing/servers/tcp"
 	"v2ray.com/core/testing/servers/udp"
-	. "v2ray.com/ext/assert"
 )
 
 func TestShadowsocksAES256TCP(t *testing.T) {
@@ -538,8 +536,6 @@ func TestShadowsocksAES256GCMTCP(t *testing.T) {
 }
 
 func TestShadowsocksAES128GCMUDP(t *testing.T) {
-	assert := With(t)
-
 	udpServer := udp.Server{
 		MsgProcessor: xor,
 	}
@@ -626,38 +622,18 @@ func TestShadowsocksAES128GCMUDP(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
-	var wg sync.WaitGroup
-	wg.Add(10)
+	var errg errgroup.Group
 	for i := 0; i < 10; i++ {
-		go func() {
-			conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
-				IP:   []byte{127, 0, 0, 1},
-				Port: int(clientPort),
-			})
-			common.Must(err)
-
-			payload := make([]byte, 1024)
-			rand.Read(payload)
-
-			nBytes, err := conn.Write([]byte(payload))
-			common.Must(err)
-			assert(nBytes, Equals, len(payload))
-
-			response := readFrom(conn, time.Second*5, 1024)
-			assert(response, Equals, xor([]byte(payload)))
-			assert(conn.Close(), IsNil)
-			wg.Done()
-		}()
+		errg.Go(testUDPConn(clientPort, 1024, time.Second*5))
+	}
+	if err := errg.Wait(); err != nil {
+		t.Error(err)
 	}
-	wg.Wait()
-
-	CloseAllServers(servers)
 }
 
 func TestShadowsocksAES128GCMUDPMux(t *testing.T) {
-	assert := With(t)
-
 	udpServer := udp.Server{
 		MsgProcessor: xor,
 	}
@@ -750,33 +726,15 @@ func TestShadowsocksAES128GCMUDPMux(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
-	var wg sync.WaitGroup
-	wg.Add(10)
+	var errg errgroup.Group
 	for i := 0; i < 10; i++ {
-		go func() {
-			conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
-				IP:   []byte{127, 0, 0, 1},
-				Port: int(clientPort),
-			})
-			common.Must(err)
-
-			payload := make([]byte, 1024)
-			rand.Read(payload)
-
-			nBytes, err := conn.Write([]byte(payload))
-			common.Must(err)
-			assert(nBytes, Equals, len(payload))
-
-			response := readFrom(conn, time.Second*5, 1024)
-			assert(response, Equals, xor([]byte(payload)))
-			assert(conn.Close(), IsNil)
-			wg.Done()
-		}()
+		errg.Go(testUDPConn(clientPort, 1024, time.Second*5))
+	}
+	if err := errg.Wait(); err != nil {
+		t.Error(err)
 	}
-	wg.Wait()
-
-	CloseAllServers(servers)
 }
 
 func TestShadowsocksNone(t *testing.T) {

testing/scenarios/socks_test.go

@@ -20,7 +20,6 @@ import (
 	"v2ray.com/core/proxy/socks"
 	"v2ray.com/core/testing/servers/tcp"
 	"v2ray.com/core/testing/servers/udp"
-	. "v2ray.com/ext/assert"
 )
 
 func TestSocksBridgeTCP(t *testing.T) {
@@ -105,8 +104,6 @@ func TestSocksBridgeTCP(t *testing.T) {
 }
 
 func TestSocksBridageUDP(t *testing.T) {
-	assert := With(t)
-
 	udpServer := udp.Server{
 		MsgProcessor: xor,
 	}
@@ -180,30 +177,14 @@ func TestSocksBridageUDP(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
-	conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
-		IP:   []byte{127, 0, 0, 1},
-		Port: int(clientPort),
-	})
-	common.Must(err)
-
-	payload := "dokodemo request."
-	nBytes, err := conn.Write([]byte(payload))
-	common.Must(err)
-	assert(nBytes, Equals, len(payload))
-
-	response := make([]byte, 1024)
-	nBytes, err = conn.Read(response)
-	common.Must(err)
-	assert(response[:nBytes], Equals, xor([]byte(payload)))
-	assert(conn.Close(), IsNil)
-
-	CloseAllServers(servers)
+	if err := testUDPConn(clientPort, 1024, time.Second*5)(); err != nil {
+		t.Error(err)
+	}
 }
 
 func TestSocksBridageUDPWithRouting(t *testing.T) {
-	assert := With(t)
-
 	udpServer := udp.Server{
 		MsgProcessor: xor,
 	}
@@ -283,30 +264,14 @@ func TestSocksBridageUDPWithRouting(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
-	conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
-		IP:   []byte{127, 0, 0, 1},
-		Port: int(clientPort),
-	})
-	common.Must(err)
-
-	payload := "dokodemo request."
-	nBytes, err := conn.Write([]byte(payload))
-	common.Must(err)
-	assert(nBytes, Equals, len(payload))
-
-	response := make([]byte, 1024)
-	nBytes, err = conn.Read(response)
-	common.Must(err)
-	assert(response[:nBytes], Equals, xor([]byte(payload)))
-	assert(conn.Close(), IsNil)
-
-	CloseAllServers(servers)
+	if err := testUDPConn(clientPort, 1024, time.Second*5)(); err != nil {
+		t.Error(err)
+	}
 }
 
 func TestSocksConformance(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -356,23 +321,18 @@ func TestSocksConformance(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
 	{
 		noAuthDialer, err := xproxy.SOCKS5("tcp", net.TCPDestination(net.LocalHostIP, noAuthPort).NetAddr(), nil, xproxy.Direct)
 		common.Must(err)
 		conn, err := noAuthDialer.Dial("tcp", dest.NetAddr())
 		common.Must(err)
+		defer conn.Close()
 
-		payload := "test payload"
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-		assert(conn.Close(), IsNil)
+		if err := testTCPConn2(conn, 1024, time.Second*5)(); err != nil {
+			t.Error(err)
+		}
 	}
 
 	{
@@ -380,52 +340,32 @@ func TestSocksConformance(t *testing.T) {
 		common.Must(err)
 		conn, err := authDialer.Dial("tcp", dest.NetAddr())
 		common.Must(err)
+		defer conn.Close()
 
-		payload := "test payload"
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-		assert(conn.Close(), IsNil)
+		if err := testTCPConn2(conn, 1024, time.Second*5)(); err != nil {
+			t.Error(err)
+		}
 	}
 
 	{
 		dialer := socks4.DialSocksProxy(socks4.SOCKS4, net.TCPDestination(net.LocalHostIP, noAuthPort).NetAddr())
 		conn, err := dialer("tcp", dest.NetAddr())
 		common.Must(err)
+		defer conn.Close()
 
-		payload := "test payload"
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-		assert(conn.Close(), IsNil)
+		if err := testTCPConn2(conn, 1024, time.Second*5)(); err != nil {
+			t.Error(err)
+		}
 	}
 
 	{
 		dialer := socks4.DialSocksProxy(socks4.SOCKS4A, net.TCPDestination(net.LocalHostIP, noAuthPort).NetAddr())
 		conn, err := dialer("tcp", net.TCPDestination(net.LocalHostDomain, tcpServer.Port).NetAddr())
 		common.Must(err)
+		defer conn.Close()
 
-		payload := "test payload"
-		nBytes, err := conn.Write([]byte(payload))
-		common.Must(err)
-		assert(nBytes, Equals, len(payload))
-
-		response := make([]byte, 1024)
-		nBytes, err = conn.Read(response)
-		common.Must(err)
-		assert(response[:nBytes], Equals, xor([]byte(payload)))
-		assert(conn.Close(), IsNil)
+		if err := testTCPConn2(conn, 1024, time.Second*5)(); err != nil {
+			t.Error(err)
+		}
 	}
-
-	CloseAllServers(servers)
 }

testing/scenarios/vmess_test.go

@@ -1,9 +1,7 @@
 package scenarios
 
 import (
-	"crypto/rand"
 	"os"
-	"sync"
 	"testing"
 	"time"
 
@@ -26,7 +24,6 @@ import (
 	"v2ray.com/core/testing/servers/udp"
 	"v2ray.com/core/transport/internet"
 	"v2ray.com/core/transport/internet/kcp"
-	. "v2ray.com/ext/assert"
 )
 
 func TestVMessDynamicPort(t *testing.T) {
@@ -364,8 +361,6 @@ func TestVMessGCMReadv(t *testing.T) {
 }
 
 func TestVMessGCMUDP(t *testing.T) {
-	assert := With(t)
-
 	udpServer := udp.Server{
 		MsgProcessor: xor,
 	}
@@ -457,43 +452,15 @@ func TestVMessGCMUDP(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
 	common.Must(err)
-
-	var wg sync.WaitGroup
-	wg.Add(10)
-	for i := 0; i < 10; i++ {
-		go func() {
-			conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
-				IP:   []byte{127, 0, 0, 1},
-				Port: int(clientPort),
-			})
-			common.Must(err)
-
-			payload := make([]byte, 1024)
-			rand.Read(payload)
-
-			nBytes, err := conn.Write([]byte(payload))
-			common.Must(err)
-			assert(nBytes, Equals, len(payload))
-
-			payload1 := make([]byte, 1024)
-			rand.Read(payload1)
-			nBytes, err = conn.Write([]byte(payload1))
-			common.Must(err)
-			assert(nBytes, Equals, len(payload1))
-
-			response := readFrom(conn, time.Second*5, 1024)
-			assert(response, Equals, xor([]byte(payload)))
-
-			response = readFrom(conn, time.Second*5, 1024)
-			assert(response, Equals, xor([]byte(payload1)))
-
-			assert(conn.Close(), IsNil)
-			wg.Done()
-		}()
-	}
-	wg.Wait()
-
 	CloseAllServers(servers)
+
+	var errg errgroup.Group
+	for i := 0; i < 10; i++ {
+		errg.Go(testTCPConn(clientPort, 1024, time.Second*5))
+	}
+	if err := errg.Wait(); err != nil {
+		t.Error(err)
+	}
 }
 
 func TestVMessChacha20(t *testing.T) {
@@ -1076,8 +1043,6 @@ func TestVMessGCMMux(t *testing.T) {
 }
 
 func TestVMessGCMMuxUDP(t *testing.T) {
-	assert := With(t)
-
 	tcpServer := tcp.Server{
 		MsgProcessor: xor,
 	}
@@ -1196,71 +1161,17 @@ func TestVMessGCMMuxUDP(t *testing.T) {
 
 	servers, err := InitializeServerConfigs(serverConfig, clientConfig)
 	common.Must(err)
+	defer CloseAllServers(servers)
 
 	for range "abcd" {
-		var wg sync.WaitGroup
-		const nConnection = 16
-		wg.Add(nConnection * 2)
-		for i := 0; i < nConnection; i++ {
-			go func() {
-				conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
-					IP:   []byte{127, 0, 0, 1},
-					Port: int(clientPort),
-				})
-				common.Must(err)
-
-				payload := make([]byte, 10240)
-				rand.Read(payload)
-
-				xorpayload := xor(payload)
-
-				nBytes, err := conn.Write(payload)
-				common.Must(err)
-				assert(nBytes, Equals, len(payload))
-
-				response := readFrom(conn, time.Second*20, 10240)
-				assert(response, Equals, xorpayload)
-				assert(conn.Close(), IsNil)
-				wg.Done()
-			}()
+		var errg errgroup.Group
+		for i := 0; i < 16; i++ {
+			errg.Go(testTCPConn(clientPort, 10240, time.Second*20))
+			errg.Go(testUDPConn(clientUDPPort, 1024, time.Second*10))
 		}
-		for i := 0; i < nConnection; i++ {
-			go func() {
-				conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
-					IP:   []byte{127, 0, 0, 1},
-					Port: int(clientUDPPort),
-				})
-				common.Must(err)
-
-				conn.SetDeadline(time.Now().Add(time.Second * 10))
-
-				payload := make([]byte, 1024)
-				rand.Read(payload)
-
-				xorpayload := xor(payload)
-
-				for j := 0; j < 2; j++ {
-					nBytes, _, err := conn.WriteMsgUDP(payload, nil, nil)
-					common.Must(err)
-					assert(nBytes, Equals, len(payload))
-				}
-
-				response := make([]byte, 1024)
-				oob := make([]byte, 16)
-				for j := 0; j < 2; j++ {
-					nBytes, _, _, _, err := conn.ReadMsgUDP(response, oob)
-					common.Must(err)
-					assert(nBytes, Equals, 1024)
-					assert(response, Equals, xorpayload)
-				}
-
-				assert(conn.Close(), IsNil)
-				wg.Done()
-			}()
+		if err := errg.Wait(); err != nil {
+			t.Error(err)
 		}
-		wg.Wait()
 		time.Sleep(time.Second)
 	}
-
-	CloseAllServers(servers)
 }