fix: correct SSE datasource complete behaviour

[endigma]

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

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.

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.

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.

[coderabbitai]

Walkthrough

Replaces channel-driven GraphQL SSE subscription handling with direct updater calls (Update, Complete, Close) and adds a channel-based test updater with synchronization helpers; updates tests to use t.Context() and event-driven Await* helpers for deterministic per-update assertions and lifecycle waits.

Changes

Cohort / File(s)	Summary
SSE handler refactor v2/pkg/engine/datasource/graphql_datasource/graphql_sse_handler.go	Removes internal subscribe goroutine and data/err channels; routes incoming SSE events and errors directly through the updater API (Update, Complete, Close); centralizes error propagation via updater.Update; simplifies StartBlocking control flow.
Test updater and tests v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go, v2/pkg/engine/datasource/graphql_datasource/graphql_sse_handler_test.go	Adds testSubscriptionUpdaterChan (channels: updates, complete, closed) and constructor newTestSubscriptionUpdaterChan; enhances existing testSubscriptionUpdater with mutexes and helper methods; replaces slice-based assertions with AwaitUpdateWithT, AwaitComplete, AwaitClose, and uses t.Context() for subscription clients.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• chore: use tagged switches to simplify the code (QF1003) #1272 — Modifies graphql_sse_handler.go around SSE error payload parsing and routing; overlaps with the handler control-flow and error propagation changes in this PR.

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The title concisely identifies the fix to the SSE datasource’s complete behavior and aligns precisely with the PR’s updates to the handler and tests, following the Conventional Commits style. It avoids vague wording or extraneous details, ensuring clarity for reviewers scanning the project history.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

✨ Finishing touches

• 📝 Generate docstrings

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 0d60605 and db9227f.

📒 Files selected for processing (3)

• v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go (3 hunks)
• v2/pkg/engine/datasource/graphql_datasource/graphql_sse_handler.go (5 hunks)
• v2/pkg/engine/datasource/graphql_datasource/graphql_sse_handler_test.go (15 hunks)

🧰 Additional context used

🧬 Code graph analysis (3)

v2/pkg/engine/datasource/graphql_datasource/graphql_sse_handler.go (1)

v2/pkg/engine/resolve/response.go (1)

• SubscriptionCloseKindNormal (62-62)

v2/pkg/engine/datasource/graphql_datasource/graphql_sse_handler_test.go (2)

v2/pkg/engine/datasource/graphql_datasource/graphql_subscription_client.go (3)

• NewGraphQLSubscriptionClient (219-275)
• WithReadTimeout (131-135)
• WithLogger (125-129)

v2/pkg/engine/resolve/context.go (1)

• Context (16-35)

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go (1)

v2/pkg/engine/resolve/response.go (1)

• SubscriptionCloseKind (56-59)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: Build and test (go 1.25 / windows-latest)
• GitHub Check: Build and test (go 1.25 / ubuntu-latest)
• GitHub Check: Build and test (go 1.25 / windows-latest)
• GitHub Check: Build and test (go 1.25 / ubuntu-latest)

🔇 Additional comments (8)

v2/pkg/engine/datasource/graphql_datasource/graphql_sse_handler_test.go (8)

93-111: LGTM! Improved context and synchronization handling.

The changes correctly align the test with the new updater-driven flow:

• Using t.Context() for the engine context while the subscription context is immediately cancelled is appropriate
• Switching to newTestSubscriptionUpdaterChan() enables event-driven testing
• Using AwaitClose correctly validates that the subscription closes when aborted

---

148-149: LGTM! Proper complete event handling aligned with PR objectives.

The changes correctly implement the SSE complete behavior:

• Server explicitly emits complete event (lines 148-149)
• Test uses channel-based updater for event-driven validation (line 164)
• Per-update validation with AwaitUpdateWithT provides immediate feedback (lines 176-182)
• AwaitComplete explicitly verifies the complete event was received (line 184)
• Explicit select synchronization is clearer than the previous assert.Eventuallyf pattern (lines 188-192)

Also applies to: 164-164, 176-192

---

235-235: LGTM! Consistent event-driven testing pattern.

The test correctly validates SSE events with explicit "next" and "complete" event types, using the same reliable synchronization patterns as the other updated tests.

Also applies to: 248-264

---

301-301: LGTM! Correct error handling with AwaitClose.

The test correctly validates error scenarios by:

• Using per-update callback to verify the error payload (lines 314-316)
• Calling AwaitClose instead of AwaitComplete since errors close the connection without normal completion (line 318)
• Proper server synchronization with explicit timeout (lines 322-326)

Also applies to: 314-326

---

404-404: LGTM! Consistent error handling pattern.

The test correctly applies the same event-driven error handling pattern used in other error scenarios.

Also applies to: 417-429

---

473-473: LGTM! Proper parameter validation with consistent synchronization.

The test correctly validates query parameter handling with the updated synchronization patterns.

Also applies to: 489-501

---

614-614: LGTM! Proper handling of upstream failure scenario.

The test correctly validates the behavior when the upstream server dies:

• First update with successful message is validated (lines 627-629)
• Second update with internal error is validated (lines 631-633)
• Connection closure is awaited since the upstream terminates abnormally (line 635)
• Proper server synchronization (lines 639-643)

Also applies to: 627-643

---

20-82: Verify if test should be updated to new pattern.

TestGraphQLSubscriptionClientSubscribe_SSE (lines 20-82) still uses the old synchronization pattern with &testSubscriptionUpdater{}, AwaitUpdates, and assert.Eventuallyf, while all other SSE tests have been updated to use newTestSubscriptionUpdaterChan(), AwaitUpdateWithT, and explicit select statements.

Is this intentional (e.g., testing a different code path or maintaining backward compatibility coverage), or should this test also be updated for consistency?

[coderabbitai]

Actionable comments posted: 0

♻️ Duplicate comments (1)

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go (1)

8316-8325: Shadowing bug in AwaitCloseKind fixed — LGTM

Received value no longer shadows the expected; assertion now compares expectedCloseKind to received.

🧹 Nitpick comments (1)

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go (1)

8265-8277: Make Complete/Close idempotent to avoid test flakiness/panics

• Complete closes an unguarded channel: a second Complete will panic.
• Close uses an unbuffered chan and is not idempotent: a duplicate Close may block if no receiver.

Recommend guarding Complete/Close with sync.Once and buffer closed chan to 1 to safely signal exactly once.

Apply:

 type testSubscriptionUpdaterChan struct {
-	updates  chan string
-	complete chan struct{}
-	closed   chan resolve.SubscriptionCloseKind
+	updates      chan string
+	complete     chan struct{}
+	closed       chan resolve.SubscriptionCloseKind
+	onceComplete sync.Once
+	onceClosed   sync.Once
 }
 
 func newTestSubscriptionUpdaterChan() *testSubscriptionUpdaterChan {
 	return &testSubscriptionUpdaterChan{
 		updates:  make(chan string),
 		complete: make(chan struct{}),
-		closed:   make(chan resolve.SubscriptionCloseKind),
+		closed:   make(chan resolve.SubscriptionCloseKind, 1),
 	}
 }
 
 func (t *testSubscriptionUpdaterChan) Complete() {
-	close(t.complete)
+	t.onceComplete.Do(func() { close(t.complete) })
 }
 
 func (t *testSubscriptionUpdaterChan) Close(kind resolve.SubscriptionCloseKind) {
-	t.closed <- kind
+	t.onceClosed.Do(func() { t.closed <- kind })
 }

Also applies to: 8287-8293, 8291-8293

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between db9227f and 3dc4534.

📒 Files selected for processing (1)

• v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go (3 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go (1)

v2/pkg/engine/resolve/response.go (1)

• SubscriptionCloseKind (56-59)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: Build and test (go 1.25 / windows-latest)
• GitHub Check: Build and test (go 1.25 / ubuntu-latest)
• GitHub Check: Build and test (go 1.25 / windows-latest)
• GitHub Check: Build and test (go 1.25 / ubuntu-latest)

🔇 Additional comments (1)

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_test.go (1)

8347-8348: Good test ergonomics: added tt.Helper()

Helper frames failures at the caller. Keep.

Also applies to: 8368-8369

[StarpTech]

LGTM