From 97475071e15a107e30485df48ca201ef73768d0e Mon Sep 17 00:00:00 2001
From: Sandy Tao <sandytao520@icloud.com>
Date: Mon, 15 Sep 2025 09:36:35 -0700
Subject: [PATCH] fix(tests): Cleans up setup test and fix behavior test
 (#8391)

---
 packages/core/src/core/client.test.ts | 48 +++++----------------------
 1 file changed, 9 insertions(+), 39 deletions(-)

diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts
index 37f002a527..5a806ced66 100644
--- a/packages/core/src/core/client.test.ts
+++ b/packages/core/src/core/client.test.ts
@@ -2320,39 +2320,6 @@ ${JSON.stringify(
       expect(mockCheckNextSpeaker).not.toHaveBeenCalled();
     });
 
-    it('should create linked abort signal and pass it to turn.run', async () => {
-      // Arrange
-      const mockStream = (async function* () {
-        yield { type: 'content', value: 'Hello' };
-      })();
-      mockTurnRunFn.mockReturnValue(mockStream);
-
-      const mockChat: Partial<GeminiChat> = {
-        addHistory: vi.fn(),
-        getHistory: vi.fn().mockReturnValue([]),
-      };
-      client['chat'] = mockChat as GeminiChat;
-
-      const originalSignal = new AbortController().signal;
-
-      // Act
-      const stream = client.sendMessageStream(
-        [{ text: 'Hi' }],
-        originalSignal,
-        'prompt-id-signal',
-      );
-      for await (const _ of stream) {
-        // consume stream
-      }
-
-      // Assert
-      expect(mockTurnRunFn).toHaveBeenCalledWith(
-        expect.any(String),
-        [{ text: 'Hi' }],
-        expect.not.objectContaining({ signal: originalSignal }),
-      );
-    });
-
     it('should abort linked signal when loop is detected', async () => {
       // Arrange
       vi.spyOn(client['loopDetector'], 'turnStarted').mockResolvedValue(false);
@@ -2360,11 +2327,14 @@ ${JSON.stringify(
         .mockReturnValueOnce(false)
         .mockReturnValueOnce(true);
 
-      const mockStream = (async function* () {
-        yield { type: 'content', value: 'First event' };
-        yield { type: 'content', value: 'Second event' };
-      })();
-      mockTurnRunFn.mockReturnValue(mockStream);
+      let capturedSignal: AbortSignal;
+      mockTurnRunFn.mockImplementation((model, request, signal) => {
+        capturedSignal = signal;
+        return (async function* () {
+          yield { type: 'content', value: 'First event' };
+          yield { type: 'content', value: 'Second event' };
+        })();
+      });
 
       const mockChat: Partial<GeminiChat> = {
         addHistory: vi.fn(),
@@ -2386,7 +2356,7 @@ ${JSON.stringify(
 
       // Assert
       expect(events).toContainEqual({ type: GeminiEventType.LoopDetected });
-      expect(client['loopDetector'].addAndCheck).toHaveBeenCalledTimes(2);
+      expect(capturedSignal!.aborted).toBe(true);
     });
   });