From f286ae536e75ccf53165357708a413857e5c8221 Mon Sep 17 00:00:00 2001
From: Yaacov Rydzinski <yaacovCR@gmail.com>
Date: Mon, 27 Apr 2026 15:56:33 +0300
Subject: [PATCH] polish: clean up stream tests

---
 .../incremental/__tests__/stream-test.ts      | 25 +++++--------------
 .../__tests__/legacy-stream-test.ts           | 25 +++++--------------
 2 files changed, 12 insertions(+), 38 deletions(-)

diff --git a/src/execution/incremental/__tests__/stream-test.ts b/src/execution/incremental/__tests__/stream-test.ts
index 9f92f92e60..a2d955fa46 100644
--- a/src/execution/incremental/__tests__/stream-test.ts
+++ b/src/execution/incremental/__tests__/stream-test.ts
@@ -3851,25 +3851,12 @@ describe('Execute: stream directive (cancellation)', () => {
     abortController.abort();
     await resolveOnNextTick();
 
-    let firstResult:
-      | IteratorResult<SubsequentIncrementalExecutionResult>
-      | undefined;
-    try {
-      firstResult =
-        (await nextResultPromise) as IteratorResult<SubsequentIncrementalExecutionResult>;
-    } catch (error) {
-      expect(error).to.be.instanceOf(Error);
-      expect((error as Error).message).to.equal('This operation was aborted');
-    }
-    if (firstResult && !firstResult.done) {
-      try {
-        const followUp = await iterator.next();
-        expect(followUp.done).to.equal(true);
-      } catch (error) {
-        expect(error).to.be.instanceOf(Error);
-        expect((error as Error).message).to.equal('This operation was aborted');
-      }
-    }
+    await expectPromise(nextResultPromise).toRejectWith(
+      'This operation was aborted',
+    );
+    await resolveOnNextTick();
+    const followUp = await iterator.next();
+    expect(followUp.done).to.equal(true);
     expect(streamReturnCount).to.equal(1);
 
     const priorStreamReturnCount = streamReturnCount;
diff --git a/src/execution/legacyIncremental/__tests__/legacy-stream-test.ts b/src/execution/legacyIncremental/__tests__/legacy-stream-test.ts
index ceaa05d587..122644db58 100644
--- a/src/execution/legacyIncremental/__tests__/legacy-stream-test.ts
+++ b/src/execution/legacyIncremental/__tests__/legacy-stream-test.ts
@@ -3640,25 +3640,12 @@ describe('Execute: stream directive (legacy cancellation)', () => {
     abortController.abort();
     await resolveOnNextTick();
 
-    let firstResult:
-      | IteratorResult<LegacySubsequentIncrementalExecutionResult>
-      | undefined;
-    try {
-      firstResult =
-        (await nextResultPromise) as IteratorResult<LegacySubsequentIncrementalExecutionResult>;
-    } catch (error) {
-      expect(error).to.be.instanceOf(Error);
-      expect((error as Error).message).to.equal('This operation was aborted');
-    }
-    if (firstResult && !firstResult.done) {
-      try {
-        const followUp = await iterator.next();
-        expect(followUp.done).to.equal(true);
-      } catch (error) {
-        expect(error).to.be.instanceOf(Error);
-        expect((error as Error).message).to.equal('This operation was aborted');
-      }
-    }
+    await expectPromise(nextResultPromise).toRejectWith(
+      'This operation was aborted',
+    );
+    await resolveOnNextTick();
+    const followUp = await iterator.next();
+    expect(followUp.done).to.equal(true);
     expect(streamReturnCount).to.equal(1);
 
     const priorStreamReturnCount = streamReturnCount;