Update tests to pass after refactor to Response

Author: erikeldridge

packages/vertexai/src/methods/generate-content.test.ts

@@ -309,24 +309,23 @@ describe('generateContent()', () => {
     expect(mockFetch).to.be.called;
   });
   it('on-device', async () => {
-    const expectedText = 'hi';
     const chromeAdapter = new ChromeAdapter();
-    const mockIsAvailable = stub(chromeAdapter, 'isAvailable').resolves(true);
-    const mockGenerateContent = stub(
+    const isAvailableStub = stub(chromeAdapter, 'isAvailable').resolves(true);
+    const mockResponse = getMockResponse(
+      'unary-success-basic-reply-short.json'
+    );
+    const generateContentStub = stub(
       chromeAdapter,
       'generateContentOnDevice'
-    ).resolves({
-      text: () => expectedText,
-      functionCalls: () => undefined
-    });
+    ).resolves(mockResponse as Response);
     const result = await generateContent(
       fakeApiSettings,
       'model',
       fakeRequestParams,
       chromeAdapter
     );
-    expect(result.response.text()).to.equal(expectedText);
-    expect(mockIsAvailable).to.be.called;
-    expect(mockGenerateContent).to.be.calledWith(fakeRequestParams);
+    expect(result.response.text()).to.include('Mountain View, California');
+    expect(isAvailableStub).to.be.called;
+    expect(generateContentStub).to.be.calledWith(fakeRequestParams);
   });
 });

packages/vertexai/src/methods/generate-content.ts

@@ -16,7 +16,6 @@
  */
 
 import {
-  EnhancedGenerateContentResponse,
   GenerateContentRequest,
   GenerateContentResponse,
   GenerateContentResult,
@@ -51,18 +50,15 @@ async function generateContentOnCloud(
   model: string,
   params: GenerateContentRequest,
   requestOptions?: RequestOptions
-): Promise<EnhancedGenerateContentResponse> {
-  const response = await makeRequest(
+): Promise<Response> {
+  return makeRequest(
     model,
     Task.GENERATE_CONTENT,
     apiSettings,
     /* stream */ false,
     JSON.stringify(params),
     requestOptions
   );
-  const responseJson: GenerateContentResponse = await response.json();
-  const enhancedResponse = createEnhancedContentResponse(responseJson);
-  return enhancedResponse;
 }
 
 export async function generateContent(
@@ -72,17 +68,19 @@ export async function generateContent(
   chromeAdapter: ChromeAdapter,
   requestOptions?: RequestOptions
 ): Promise<GenerateContentResult> {
-  let enhancedResponse;
+  let response;
   if (await chromeAdapter.isAvailable(params)) {
-    enhancedResponse = await chromeAdapter.generateContentOnDevice(params);
+    response = await chromeAdapter.generateContentOnDevice(params);
   } else {
-    enhancedResponse = await generateContentOnCloud(
+    response = await generateContentOnCloud(
       apiSettings,
       model,
       params,
       requestOptions
     );
   }
+  const responseJson: GenerateContentResponse = await response.json();
+  const enhancedResponse = createEnhancedContentResponse(responseJson);
   return {
     response: enhancedResponse
   };