fix some tests

pion · Jul 19, 2024 · 03c8502 · 03c8502
1 parent b627914
commit 03c8502
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 8 deletions.
diff --git a/conn_go_test.go b/conn_go_test.go
@@ -64,9 +64,15 @@ func TestContextConfig(t *testing.T) {
 	}{
 		"Dial": {
 			f: func() (func() (net.Conn, error), func()) {
+				ctx, cancel := context.WithTimeout(context.Background(), 40*time.Millisecond)
 				return func() (net.Conn, error) {
-						return Dial("udp", addr, config)
+						conn, err := Dial("udp", addr, config)
+						if err != nil {
+							return nil, err
+						}
+						return conn, conn.HandshakeContext(ctx)
 					}, func() {
+						cancel()
 					}
 			},
 			order: []byte{0, 1, 2},
@@ -75,7 +81,11 @@ func TestContextConfig(t *testing.T) {
 			f: func() (func() (net.Conn, error), func()) {
 				ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
 				return func() (net.Conn, error) {
-						return DialWithContext(ctx, "udp", addr, config)
+						conn, err := DialWithContext(ctx, "udp", addr, config)
+						if err != nil {
+							return nil, err
+						}
+						return conn, conn.HandshakeContext(ctx)
 					}, func() {
 						cancel()
 					}
@@ -85,10 +95,16 @@ func TestContextConfig(t *testing.T) {
 		"Client": {
 			f: func() (func() (net.Conn, error), func()) {
 				ca, _ := dpipe.Pipe()
+				ctx, cancel := context.WithTimeout(context.Background(), 40*time.Millisecond)
 				return func() (net.Conn, error) {
-						return Client(dtlsnet.PacketConnFromConn(ca), ca.RemoteAddr(), config)
+						conn, err := Client(dtlsnet.PacketConnFromConn(ca), ca.RemoteAddr(), config)
+						if err != nil {
+							return nil, err
+						}
+						return conn, conn.HandshakeContext(ctx)
 					}, func() {
 						_ = ca.Close()
+						cancel()
 					}
 			},
 			order: []byte{0, 1, 2},
@@ -98,7 +114,11 @@ func TestContextConfig(t *testing.T) {
 				ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
 				ca, _ := dpipe.Pipe()
 				return func() (net.Conn, error) {
-						return ClientWithContext(ctx, dtlsnet.PacketConnFromConn(ca), ca.RemoteAddr(), config)
+						conn, err := ClientWithContext(ctx, dtlsnet.PacketConnFromConn(ca), ca.RemoteAddr(), config)
+						if err != nil {
+							return nil, err
+						}
+						return conn, conn.HandshakeContext(ctx)
 					}, func() {
 						cancel()
 						_ = ca.Close()
@@ -109,10 +129,16 @@ func TestContextConfig(t *testing.T) {
 		"Server": {
 			f: func() (func() (net.Conn, error), func()) {
 				ca, _ := dpipe.Pipe()
+				ctx, cancel := context.WithTimeout(context.Background(), 40*time.Millisecond)
 				return func() (net.Conn, error) {
-						return Server(dtlsnet.PacketConnFromConn(ca), ca.RemoteAddr(), config)
+						conn, err := Server(dtlsnet.PacketConnFromConn(ca), ca.RemoteAddr(), config)
+						if err != nil {
+							return nil, err
+						}
+						return conn, conn.HandshakeContext(ctx)
 					}, func() {
 						_ = ca.Close()
+						cancel()
 					}
 			},
 			order: []byte{0, 1, 2},
@@ -122,7 +148,11 @@ func TestContextConfig(t *testing.T) {
 				ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
 				ca, _ := dpipe.Pipe()
 				return func() (net.Conn, error) {
-						return ServerWithContext(ctx, dtlsnet.PacketConnFromConn(ca), ca.RemoteAddr(), config)
+						conn, err := ServerWithContext(ctx, dtlsnet.PacketConnFromConn(ca), ca.RemoteAddr(), config)
+						if err != nil {
+							return nil, err
+						}
+						return conn, conn.HandshakeContext(ctx)
 					}, func() {
 						cancel()
 						_ = ca.Close()

diff --git a/conn_test.go b/conn_test.go
@@ -295,7 +295,11 @@ func testClient(ctx context.Context, c net.PacketConn, rAddr net.Addr, cfg *Conf
 		cfg.Certificates = []tls.Certificate{clientCert}
 	}
 	cfg.InsecureSkipVerify = true
-	return ClientWithContext(ctx, c, rAddr, cfg)
+	conn, err := ClientWithContext(ctx, c, rAddr, cfg)
+	if err != nil {
+		return nil, err
+	}
+	return conn, conn.HandshakeContext(ctx)
 }
 
 func testServer(ctx context.Context, c net.PacketConn, rAddr net.Addr, cfg *Config, generateCertificate bool) (*Conn, error) {
@@ -306,7 +310,11 @@ func testServer(ctx context.Context, c net.PacketConn, rAddr net.Addr, cfg *Conf
 		}
 		cfg.Certificates = []tls.Certificate{serverCert}
 	}
-	return ServerWithContext(ctx, c, rAddr, cfg)
+	conn, err := ServerWithContext(ctx, c, rAddr, cfg)
+	if err != nil {
+		return nil, err
+	}
+	return conn, conn.HandshakeContext(ctx)
 }
 
 func sendClientHello(cookie []byte, ca net.Conn, sequenceNumber uint64, extensions []extension.Extension) error {