fix: correct SSE datasource complete behaviour (#1311)

- **fix: clean up SSE handler and properly handle complete**
- **test: new testSubscriptionUpdaterChan**
- **test: update SSE datasource tests to assert correct complete
behaviour**

<!--
Important: Before developing new features, please open an issue to
discuss your ideas with the maintainers. This ensures project alignment
and helps avoid unnecessary work for you.

Thank you for your contribution! Please provide a detailed description
below and ensure you've met all the requirements.

Squashed commit messages must follow the [Conventional
Commits](https://www.conventionalcommits.org/en/v1.0.0/) standard to
facilitate changelog generation.

Please ensure your PR title follows the Conventional Commits
specification, using the appropriate type (e.g., feat, fix, docs) and
scope.

Examples of good PR titles:

- 💥feat!: change implementation in an non-backward compatible way
- ✨feat(auth): add support for OAuth2 login
- 🐞fix(router): add support for custom metrics
- 📚docs(README): update installation instructions
- 🧹chore(deps): bump dependencies to latest versions
-->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- Refactor
- Reworked SSE subscription handling to process events synchronously and
centralize error and completion propagation for more deterministic
behavior.
- Bug Fixes
- Reduced race conditions and flaky behavior in subscription flows via
stronger synchronization and guarded state transitions.
- Tests
- Added event-driven test helpers and updated tests to assert per-event
updates, completion, and closure with timeout-based coordination for
more reliable coverage.
- Documentation
- Added notes clarifying test lifecycle and deprecation guidance for
older test utilities.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

## Checklist

- [ ] I have discussed my proposed changes in an issue and have received
approval to proceed.
- [ ] I have followed the coding standards of the project.
- [ ] Tests or benchmarks have been added or updated.

<!--
Please add any additional information or context regarding your changes
here.
-->

#### Bug
Complete events from SSE subgraph subscriptions were not triggering the
updater correctly, this lead to hanged subscriptions because complete
was not causing client disconnects. Our test suite had poor coverage for
this behaviour.

#### Fix
- Simplified the SSE datasource (by-product of debugging, and it would
have been more complex and less clear what was happening and why if it
was "hotfixed")
- Correctly trigger updater complete when the complete event is sent
from the subgraph
- Complete events trigger closure, and for every other case we have a
deferred close.
- No more leaking goroutines or stuck connections

#### Other
- Made a new testing subscription updater to make writing the tests
easier and fix some existing flawed checks that could result in a
timeout or flaky failure in the future.

---------

Co-authored-by: Alessandro Pagnin <ale@wundergraph.com>

Author: endigma

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go

@@ -8262,6 +8262,80 @@ func (f *FailingSubscriptionClient) UniqueRequestID(ctx *resolve.Context, option
 	return errSubscriptionClientFail
 }
 
+type testSubscriptionUpdaterChan struct {
+	updates  chan string
+	complete chan struct{}
+	closed   chan resolve.SubscriptionCloseKind
+}
+
+func newTestSubscriptionUpdaterChan() *testSubscriptionUpdaterChan {
+	return &testSubscriptionUpdaterChan{
+		updates:  make(chan string),
+		complete: make(chan struct{}),
+		closed:   make(chan resolve.SubscriptionCloseKind),
+	}
+}
+
+func (t *testSubscriptionUpdaterChan) Heartbeat() {
+	t.updates <- "{}"
+}
+
+func (t *testSubscriptionUpdaterChan) Update(data []byte) {
+	t.updates <- string(data)
+}
+
+func (t *testSubscriptionUpdaterChan) Complete() {
+	close(t.complete)
+}
+
+func (t *testSubscriptionUpdaterChan) Close(kind resolve.SubscriptionCloseKind) {
+	t.closed <- kind
+}
+
+func (t *testSubscriptionUpdaterChan) AwaitUpdateWithT(tt *testing.T, timeout time.Duration, f func(t *testing.T, update string), msgAndArgs ...any) {
+	tt.Helper()
+
+	select {
+	case args := <-t.updates:
+		f(tt, args)
+	case <-time.After(timeout):
+		require.Fail(tt, "unable to receive update before timeout", msgAndArgs...)
+	}
+}
+
+func (t *testSubscriptionUpdaterChan) AwaitClose(tt *testing.T, timeout time.Duration, msgAndArgs ...any) {
+	tt.Helper()
+
+	select {
+	case <-t.closed:
+	case <-time.After(timeout):
+		require.Fail(tt, "updater not closed before timeout", msgAndArgs...)
+	}
+}
+
+func (t *testSubscriptionUpdaterChan) AwaitCloseKind(tt *testing.T, timeout time.Duration, expectedCloseKind resolve.SubscriptionCloseKind, msgAndArgs ...any) {
+	tt.Helper()
+
+	select {
+	case closeKind := <-t.closed:
+		require.Equal(tt, expectedCloseKind, closeKind, msgAndArgs...)
+	case <-time.After(timeout):
+		require.Fail(tt, "updater not closed before timeout", msgAndArgs...)
+	}
+}
+
+func (t *testSubscriptionUpdaterChan) AwaitComplete(tt *testing.T, timeout time.Duration, msgAndArgs ...any) {
+	tt.Helper()
+
+	select {
+	case <-t.complete:
+	case <-time.After(timeout):
+		require.Fail(tt, "updater not completed before timeout", msgAndArgs...)
+	}
+}
+
+// !! If you see this in a test you're working on, please replace it with the new testSubscriptionUpdaterChan
+// It's faster, more ergonomic and more reliable. See SSE handler tests for usage examples.
 type testSubscriptionUpdater struct {
 	updates []string
 	done    bool
@@ -8270,6 +8344,8 @@ type testSubscriptionUpdater struct {
 }
 
 func (t *testSubscriptionUpdater) AwaitUpdates(tt *testing.T, timeout time.Duration, count int) {
+	tt.Helper()
+
 	ticker := time.NewTicker(timeout)
 	defer ticker.Stop()
 	for {
@@ -8289,6 +8365,8 @@ func (t *testSubscriptionUpdater) AwaitUpdates(tt *testing.T, timeout time.Durat
 }
 
 func (t *testSubscriptionUpdater) AwaitDone(tt *testing.T, timeout time.Duration) {
+	tt.Helper()
+
 	ticker := time.NewTicker(timeout)
 	defer ticker.Stop()
 	for {

v2/pkg/engine/datasource/graphql_datasource/graphql_sse_handler.go

@@ -44,32 +44,8 @@ func newSSEConnectionHandler(requestContext, engineContext context.Context, conn
 }
 
 func (h *gqlSSEConnectionHandler) StartBlocking() {
-	dataCh := make(chan []byte)
-	errCh := make(chan []byte)
-	defer func() {
-		close(dataCh)
-		close(errCh)
-		h.updater.Complete()
-	}()
-
-	go h.subscribe(dataCh, errCh)
+	defer h.updater.Close(resolve.SubscriptionCloseKindNormal)
 
-	for {
-		select {
-		case data := <-dataCh:
-			h.updater.Update(data)
-		case data := <-errCh:
-			h.updater.Update(data)
-			return
-		case <-h.requestContext.Done():
-			return
-		case <-h.engineContext.Done():
-			return
-		}
-	}
-}
-
-func (h *gqlSSEConnectionHandler) subscribe(dataCh, errCh chan []byte) {
 	resp, err := h.performSubscriptionRequest()
 	if err != nil {
 		h.log.Error("failed to perform subscription request", log.Error(err))
@@ -83,6 +59,7 @@ func (h *gqlSSEConnectionHandler) subscribe(dataCh, errCh chan []byte) {
 
 		return
 	}
+
 	defer func() {
 		_ = resp.Body.Close()
 	}()
@@ -105,8 +82,7 @@ func (h *gqlSSEConnectionHandler) subscribe(dataCh, errCh chan []byte) {
 			}
 
 			h.log.Error("failed to read event", log.Error(err))
-
-			errCh <- []byte(internalError)
+			h.updater.Update([]byte(internalError))
 			return
 		}
 
@@ -131,12 +107,13 @@ func (h *gqlSSEConnectionHandler) subscribe(dataCh, errCh chan []byte) {
 					return
 				}
 
-				dataCh <- data
+				h.updater.Update(data)
 			case bytes.HasPrefix(line, headerEvent):
 				event := trim(line[len(headerEvent):])
 
 				switch {
 				case bytes.Equal(event, eventTypeComplete):
+					h.updater.Complete()
 					return
 				case bytes.Equal(event, eventTypeNext):
 					continue
@@ -165,33 +142,31 @@ func (h *gqlSSEConnectionHandler) subscribe(dataCh, errCh chan []byte) {
 						response, err = jsonparser.Set(response, val, "errors")
 						if err != nil {
 							h.log.Error("failed to set errors", log.Error(err))
-
-							errCh <- []byte(internalError)
+							h.updater.Update([]byte(internalError))
 							return
 						}
-						errCh <- response
+						h.updater.Update(response)
 						return
 					case jsonparser.Object:
 						response := []byte(`{"errors":[]}`)
 						response, err = jsonparser.Set(response, val, "errors", "[0]")
 						if err != nil {
 							h.log.Error("failed to set errors", log.Error(err))
-
-							errCh <- []byte(internalError)
+							h.updater.Update([]byte(internalError))
 							return
 						}
-						errCh <- response
+						h.updater.Update(response)
 						return
 					default:
 						// don't crash on unexpected payloads from upstream
 						h.log.Error(fmt.Sprintf("unexpected value type: %d", valueType))
-						errCh <- []byte(internalError)
+						h.updater.Update([]byte(internalError))
 						return
 					}
 
 				default:
 					h.log.Error("failed to parse errors", log.Error(err))
-					errCh <- []byte(internalError)
+					h.updater.Update([]byte(internalError))
 					return
 				}
 			}
@@ -210,7 +185,6 @@ func trim(data []byte) []byte {
 }
 
 func (h *gqlSSEConnectionHandler) performSubscriptionRequest() (*http.Response, error) {
-
 	var req *http.Request
 	var err error