From f20f7797a3289d8536ff826b97f1ccb223ced01c Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Fri, 13 Mar 2026 20:02:50 +0800
Subject: [PATCH] refactor(tunnel): improve error handling in gRPC stream tests

- Updated test cases to handle potential stream closure more gracefully by checking for send errors before proceeding with receive operations.
- Enhanced clarity in test logic to ensure that expected errors are correctly asserted when metadata is missing or channel IDs are non-existent.

Made-with: Cursor
---
 tai/tunnel/grpc_handler_test.go | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/tai/tunnel/grpc_handler_test.go b/tai/tunnel/grpc_handler_test.go
index 4246aa6b..278dd300 100644
--- a/tai/tunnel/grpc_handler_test.go
+++ b/tai/tunnel/grpc_handler_test.go
@@ -204,13 +204,16 @@ func TestForward_MissingMetadata(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	err = stream.Send(&taipb.ForwardData{Data: []byte("hello")})
-	if err != nil {
-		t.Fatal(err)
+
+	// Server may close the stream before or after Send completes (race).
+	// Either Send or Recv returning an error confirms the server rejected.
+	sendErr := stream.Send(&taipb.ForwardData{Data: []byte("hello")})
+	if sendErr != nil {
+		return // server already closed stream — pass
 	}
 
-	_, err = stream.Recv()
-	if err == nil {
+	_, recvErr := stream.Recv()
+	if recvErr == nil {
 		t.Fatal("expected error for missing channel_id metadata")
 	}
 }
@@ -224,13 +227,14 @@ func TestForward_NoPendingChannel(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	err = stream.Send(&taipb.ForwardData{Data: []byte("hello")})
-	if err != nil {
-		t.Fatal(err)
+
+	sendErr := stream.Send(&taipb.ForwardData{Data: []byte("hello")})
+	if sendErr != nil {
+		return // server already closed stream — pass
 	}
 
-	_, err = stream.Recv()
-	if err == nil {
+	_, recvErr := stream.Recv()
+	if recvErr == nil {
 		t.Fatal("expected error for non-existent channel_id")
 	}
 }