From b4506a98788dba4c0b68af2f8d6e8307ebd65fd2 Mon Sep 17 00:00:00 2001
From: Marek Siarkowicz <siarkowicz@google.com>
Date: Sun, 1 Dec 2024 16:52:22 +0100
Subject: [PATCH] Expect watch fragment test to allways pass as responses
 should always fit under size

Signed-off-by: Marek Siarkowicz <siarkowicz@google.com>
---
 .../clientv3/watch_fragment_test.go           | 50 +++++--------------
 1 file changed, 13 insertions(+), 37 deletions(-)

diff --git a/tests/integration/clientv3/watch_fragment_test.go b/tests/integration/clientv3/watch_fragment_test.go
index 81450f5f9aa..a3c6f87d501 100644
--- a/tests/integration/clientv3/watch_fragment_test.go
+++ b/tests/integration/clientv3/watch_fragment_test.go
@@ -28,34 +28,18 @@ import (
 	integration2 "go.etcd.io/etcd/tests/v3/framework/integration"
 )
 
-// TestWatchFragmentDisable ensures that large watch
-// response exceeding server-side request limit can
-// arrive even without watch response fragmentation.
 func TestWatchFragmentDisable(t *testing.T) {
 	testWatchFragment(t, false, false)
 }
 
-// TestWatchFragmentDisableWithGRPCLimit verifies
-// large watch response exceeding server-side request
-// limit and client-side gRPC response receive limit
-// cannot arrive without watch events fragmentation,
-// because multiple events exceed client-side gRPC
-// response receive limit.
 func TestWatchFragmentDisableWithGRPCLimit(t *testing.T) {
 	testWatchFragment(t, false, true)
 }
 
-// TestWatchFragmentEnable ensures that large watch
-// response exceeding server-side request limit arrive
-// with watch response fragmentation.
 func TestWatchFragmentEnable(t *testing.T) {
 	testWatchFragment(t, true, false)
 }
 
-// TestWatchFragmentEnableWithGRPCLimit verifies
-// large watch response exceeding server-side request
-// limit and client-side gRPC response receive limit
-// can arrive only when watch events are fragmented.
 func TestWatchFragmentEnableWithGRPCLimit(t *testing.T) {
 	testWatchFragment(t, true, true)
 }
@@ -99,29 +83,21 @@ func testWatchFragment(t *testing.T, fragment, exceedRecvLimit bool) {
 	wch := cli.Watch(context.TODO(), "foo", opts...)
 
 	// expect 10 MiB watch response
-	select {
-	case ws := <-wch:
-		// without fragment, should exceed gRPC client receive limit
-		if !fragment && exceedRecvLimit {
-			if len(ws.Events) != 0 {
-				t.Fatalf("expected 0 events with watch fragmentation, got %d", len(ws.Events))
+	eventCount := 0
+	for eventCount < 10 {
+		select {
+		case ws := <-wch:
+			// still expect merged watch events
+			if ws.Err() != nil {
+				t.Fatalf("unexpected error %v", ws.Err())
 			}
-			exp := "code = ResourceExhausted desc = grpc: received message larger than max ("
-			if !strings.Contains(ws.Err().Error(), exp) {
-				t.Fatalf("expected 'ResourceExhausted' error, got %v", ws.Err())
-			}
-			return
-		}
+			eventCount += len(ws.Events)
 
-		// still expect merged watch events
-		if len(ws.Events) != 10 {
-			t.Fatalf("expected 10 events with watch fragmentation, got %d", len(ws.Events))
-		}
-		if ws.Err() != nil {
-			t.Fatalf("unexpected error %v", ws.Err())
+		case <-time.After(testutil.RequestTimeout):
+			t.Fatalf("took too long to receive events")
 		}
-
-	case <-time.After(testutil.RequestTimeout):
-		t.Fatalf("took too long to receive events")
+	}
+	if eventCount != 10 {
+		t.Fatalf("expected 10 events with watch fragmentation, got %d", eventCount)
 	}
 }