diff --git a/src/index.test.tsx b/src/index.test.tsx
index e78f4ff..b5529ad 100644
--- a/src/index.test.tsx
+++ b/src/index.test.tsx
@@ -6,7 +6,7 @@ import * as useChatInterface from "./useChat";
 
 const MIN_PROPS = {
   customerId: "mock-customer-id",
-  corpusIds: ["1", "2"],
+  corpusKey: "1",
   apiKey: "mock-api-key"
 };
 
diff --git a/src/useChat.test.ts b/src/useChat.test.ts
index b5f0434..82de09c 100644
--- a/src/useChat.test.ts
+++ b/src/useChat.test.ts
@@ -1,7 +1,6 @@
 import { useChat } from "./useChat";
 import { act, renderHook } from "@testing-library/react-hooks";
 import { waitFor } from "@testing-library/react";
-import * as sendSearchRequestInterface from "./utils/sendSearchRequest";
 import * as streamQueryInterface from "@vectara/stream-query-client";
 
 jest.mock("@vectara/stream-query-client", () => {
@@ -11,43 +10,11 @@ jest.mock("@vectara/stream-query-client", () => {
   };
 });
 
-const MOCK_API_RESPONSE = {
-  document: [
-    {
-      id: "mock-doc-id",
-      metadata: [{ name: "mock-name", value: "mock-value" }]
-    }
-  ],
-  response: [
-    {
-      corpusKey: {
-        customerId: 0,
-        corpusId: 1
-      },
-      documentIndex: 0,
-      metadata: [],
-      score: 0.8,
-      text: "mock-text"
-    }
-  ],
-  summary: [
-    {
-      chat: {
-        conversationId: "mock-conversation-id",
-        turnId: "mock-turn-id",
-        text: "mock-answer"
-      }
-    }
-  ]
-};
-
 describe("useChat", () => {
-  let sendSearchRequestSpy: jest.SpyInstance;
   let streamQuerySpy: jest.SpyInstance;
 
   beforeEach(() => {
-    sendSearchRequestSpy = jest.spyOn(sendSearchRequestInterface, "sendSearchRequest");
-    streamQuerySpy = jest.spyOn(streamQueryInterface, "streamQuery");
+    streamQuerySpy = jest.spyOn(streamQueryInterface, "streamQueryV2");
   });
 
   afterEach(() => {
@@ -57,16 +24,15 @@ describe("useChat", () => {
   describe("streaming", () => {
     it("should send messages and update hook values", async () => {
       const { result } = renderHook(() =>
-        useChat({ customerId: "mock-customer-id", corpusIds: ["1"], apiKey: "mock-api-key" })
+        useChat({ customerId: "mock-customer-id", corpusKey: "1", apiKey: "mock-api-key" })
       );
 
-      streamQuerySpy.mockImplementation(async (_, onStreamUpdate) => {
-        await onStreamUpdate({
-          references: [],
-          detail: null,
+      streamQuerySpy.mockImplementation(async (_, onStreamEvent) => {
+        const event = {
+          type: "generationChunk",
           updatedText: "mock-updated-text",
-          isDone: false
-        });
+        };
+        await onStreamEvent(event);
       });
 
       await act(async () => {
@@ -101,108 +67,4 @@ describe("useChat", () => {
       });
     });
   });
-
-  describe("non-streaming", () => {
-    it("should send messages and update message history", async () => {
-      const { result } = renderHook(() =>
-        useChat({ customerId: "mock-customer-id", corpusIds: ["1"], apiKey: "mock-api-key", enableStreaming: false })
-      );
-
-      sendSearchRequestSpy.mockImplementation(() => Promise.resolve(MOCK_API_RESPONSE));
-
-      await act(async () => {
-        await result.current.sendMessage({ query: "mock-query" });
-      });
-
-      expect(sendSearchRequestSpy).toHaveBeenCalledWith(
-        expect.objectContaining({
-          queryValue: "mock-query"
-        })
-      );
-
-      expect(result.current.messageHistory.length).toEqual(1);
-    });
-
-    it("should reflect error state", async () => {
-      const { result } = renderHook(() =>
-        useChat({ customerId: "mock-customer-id", corpusIds: ["1"], apiKey: "mock-api-key", enableStreaming: false })
-      );
-      sendSearchRequestSpy.mockImplementation(() => {
-        throw "error";
-      });
-
-      await act(async () => {
-        result.current.sendMessage({ query: "mock-query" });
-      });
-
-      expect(result.current.hasError).toEqual(true);
-    });
-
-    it("should reflect loading state", async () => {
-      const { result } = renderHook(() =>
-        useChat({ customerId: "mock-customer-id", corpusIds: ["1"], apiKey: "mock-api-key" })
-      );
-      sendSearchRequestSpy.mockImplementation(() => {
-        return new Promise(() => {});
-      });
-
-      act(() => {
-        result.current.sendMessage({ query: "mock-query" });
-      });
-
-      expect(result.current.isLoading).toEqual(true);
-    });
-  });
-
-  it("should be able to reset the conversation", async () => {
-    const { result } = renderHook(() =>
-      useChat({ customerId: "mock-customer-id", corpusIds: ["1"], apiKey: "mock-api-key", enableStreaming: false })
-    );
-    sendSearchRequestSpy.mockImplementation(() => Promise.resolve(MOCK_API_RESPONSE));
-
-    await act(async () => {
-      await result.current.sendMessage({ query: "mock-query" });
-      await result.current.sendMessage({ query: "mock-query-2" });
-    });
-
-    // Assert that the second request uses the current conversation id.
-    expect(sendSearchRequestSpy).toHaveBeenCalledWith(
-      expect.objectContaining({
-        chat: {
-          conversationId: "mock-conversation-id"
-        }
-      })
-    );
-
-    sendSearchRequestSpy.mockImplementation(() =>
-      Promise.resolve({
-        ...MOCK_API_RESPONSE,
-        summary: [
-          {
-            chat: {
-              conversationId: "mock-conversation-id-2",
-              turnId: "mock-turn-id",
-              text: "mock-answer"
-            }
-          }
-        ]
-      })
-    );
-
-    await act(async () => {
-      await result.current.startNewConversation();
-    });
-
-    expect(result.current.messageHistory.length).toEqual(0);
-
-    await act(async () => {
-      await result.current.sendMessage({ query: "mock-query-3" });
-    });
-
-    const calls = sendSearchRequestSpy.mock.calls;
-    const recentSendSearchRequestCall = calls[calls.length - 1][0];
-
-    // Assert that the request after reset is has no conversation id.
-    expect(recentSendSearchRequestCall.chat.conversationId).toEqual(undefined);
-  });
 });