fix(blob): fix blobsub subscription cancellation test (#3746)

blob/service_test.go

@@ -683,23 +683,28 @@ func TestService_Subscribe(t *testing.T) {
 
 	t.Run("subscription cancellation", func(t *testing.T) {
 		ns := blobs[0].Namespace()
-		subCtx, subCancel := context.WithCancel(ctx)
+		subCtx, subCancel := context.WithTimeout(ctx, time.Second*2)
 		subCh, err := service.Subscribe(subCtx, ns)
 		require.NoError(t, err)
 
 		// cancel the subscription context after receiving the first response
 		select {
 		case <-subCh:
 			subCancel()
-		case <-time.After(time.Second * 2):
+		case <-ctx.Done():
 			t.Fatal("timeout waiting for first subscription response")
 		}
 
-		select {
-		case _, ok := <-subCh:
-			assert.False(t, ok, "expected subscription channel to be closed")
-		case <-time.After(time.Second * 2):
-			t.Fatal("timeout waiting for subscription channel to close")
+		for {
+			select {
+			case _, ok := <-subCh:
+				if !ok {
+					// channel closed as expected
+					return
+				}
+			case <-ctx.Done():
+				t.Fatal("timeout waiting for subscription channel to close")
+			}
 		}
 	})