chore: make unit test more stable

Author: joway

connection_test.go

@@ -499,18 +499,15 @@ func TestConnDetach(t *testing.T) {
 func TestParallelShortConnection(t *testing.T) {
 	ln, err := createTestListener("tcp", ":12345")
 	MustNil(t, err)
-	defer ln.Close()
-
 	var received int64
 	el, err := NewEventLoop(func(ctx context.Context, connection Connection) error {
 		data, err := connection.Reader().Next(connection.Reader().Len())
-		if err != nil {
-			return err
-		}
+		Assert(t, err == nil || errors.Is(err, ErrEOF))
 		atomic.AddInt64(&received, int64(len(data)))
-		//t.Logf("conn[%s] received: %d, active: %v", connection.RemoteAddr(), len(data), connection.IsActive())
+		t.Logf("conn[%s] received: %d, active: %v", connection.RemoteAddr(), len(data), connection.IsActive())
 		return nil
 	})
+	defer el.Shutdown(context.Background())
 	go func() {
 		el.Serve(ln)
 	}()
@@ -536,10 +533,11 @@ func TestParallelShortConnection(t *testing.T) {
 	}
 	wg.Wait()
 
-	for atomic.LoadInt64(&received) < int64(totalSize) {
-		t.Logf("received: %d, except: %d", atomic.LoadInt64(&received), totalSize)
+	start := time.Now()
+	for atomic.LoadInt64(&received) < int64(totalSize) && time.Now().Sub(start) < time.Second {
 		time.Sleep(time.Millisecond * 100)
 	}
+	Equal(t, atomic.LoadInt64(&received), int64(totalSize))
 }
 
 func TestConnectionServerClose(t *testing.T) {
@@ -643,8 +641,6 @@ func TestConnectionServerClose(t *testing.T) {
 func TestConnectionDailTimeoutAndClose(t *testing.T) {
 	ln, err := createTestListener("tcp", ":12345")
 	MustNil(t, err)
-	defer ln.Close()
-
 	el, err := NewEventLoop(
 		func(ctx context.Context, connection Connection) error {
 			_, err = connection.Reader().Next(connection.Reader().Len())
@@ -668,7 +664,7 @@ func TestConnectionDailTimeoutAndClose(t *testing.T) {
 			go func() {
 				defer wg.Done()
 				conn, err := DialConnection("tcp", ":12345", time.Nanosecond)
-				Assert(t, err == nil || strings.Contains(err.Error(), "i/o timeout"))
+				Assert(t, err == nil || strings.Contains(err.Error(), "i/o timeout"), err)
 				_ = conn
 			}()
 		}

mux/shard_queue_test.go

@@ -19,35 +19,38 @@ package mux
 
 import (
 	"net"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/cloudwego/netpoll"
 )
 
 func TestShardQueue(t *testing.T) {
-	var svrConn net.Conn
 	accepted := make(chan struct{})
 
 	network, address := "tcp", ":18888"
 	ln, err := net.Listen("tcp", ":18888")
 	MustNil(t, err)
 	stop := make(chan int, 1)
-	defer close(stop)
+	count, pkgsize := 16, 11
+	var wg sync.WaitGroup
+	wg.Add(1)
 	go func() {
-		var err error
-		for {
-			select {
-			case <-stop:
-				err = ln.Close()
-				MustNil(t, err)
-				return
-			default:
-			}
-			svrConn, err = ln.Accept()
-			MustNil(t, err)
-			accepted <- struct{}{}
-		}
+		defer wg.Done()
+		svrConn, err := ln.Accept()
+		MustNil(t, err)
+		accepted <- struct{}{}
+
+		total := count * pkgsize
+		recv := make([]byte, total)
+		rn, err := svrConn.Read(recv)
+		MustNil(t, err)
+		Equal(t, rn, total)
+
+		<-stop
+		err = ln.Close()
+		MustNil(t, err)
 	}()
 
 	conn, err := netpoll.DialConnection(network, address, time.Second)
@@ -56,8 +59,7 @@ func TestShardQueue(t *testing.T) {
 
 	// test
 	queue := NewShardQueue(4, conn)
-	count, pkgsize := 16, 11
-	for i := 0; i < int(count); i++ {
+	for i := 0; i < count; i++ {
 		var getter WriterGetter = func() (buf netpoll.Writer, isNil bool) {
 			buf = netpoll.NewLinkBuffer(pkgsize)
 			buf.Malloc(pkgsize)
@@ -68,14 +70,7 @@ func TestShardQueue(t *testing.T) {
 
 	err = queue.Close()
 	MustNil(t, err)
-	total := count * pkgsize
-	recv := make([]byte, total)
-	rn, err := svrConn.Read(recv)
-	MustNil(t, err)
-	Equal(t, rn, total)
-}
 
-// TODO: need mock flush
-func BenchmarkShardQueue(b *testing.B) {
-	b.Skip()
+	close(stop)
+	wg.Wait()
 }