From d3f23e0528dd6dec8a5d170cb711f3f909f4f4ce Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 22 Nov 2024 15:40:56 -0800
Subject: [PATCH] add working vertex jest tests

---
 tests/pass_through_tests/test_local_vertex.js | 14 ++++
 tests/pass_through_tests/test_vertex.js       | 40 -----------
 tests/pass_through_tests/test_vertex.test.js  | 71 +++++++++++++++++++
 3 files changed, 85 insertions(+), 40 deletions(-)
 delete mode 100644 tests/pass_through_tests/test_vertex.js
 create mode 100644 tests/pass_through_tests/test_vertex.test.js

diff --git a/tests/pass_through_tests/test_local_vertex.js b/tests/pass_through_tests/test_local_vertex.js
index 1f2eaea33..6e8fbeacd 100644
--- a/tests/pass_through_tests/test_local_vertex.js
+++ b/tests/pass_through_tests/test_local_vertex.js
@@ -1,5 +1,19 @@
 const { VertexAI, RequestOptions } = require('@google-cloud/vertexai');
 
+
+// Import fetch if the SDK uses it
+const originalFetch = global.fetch || require('node-fetch');
+
+// Monkey-patch the fetch used internally
+global.fetch = async function patchedFetch(url, options) {
+    // Modify the URL to use HTTP instead of HTTPS
+    if (url.startsWith('https://localhost:4000')) {
+        url = url.replace('https://', 'http://');
+    }
+    console.log('Patched fetch sending request to:', url);
+    return originalFetch(url, options);
+};
+
 const vertexAI = new VertexAI({
     project: 'adroit-crow-413218',
     location: 'us-central1',
diff --git a/tests/pass_through_tests/test_vertex.js b/tests/pass_through_tests/test_vertex.js
deleted file mode 100644
index 1f2eaea33..000000000
--- a/tests/pass_through_tests/test_vertex.js
+++ /dev/null
@@ -1,40 +0,0 @@
-const { VertexAI, RequestOptions } = require('@google-cloud/vertexai');
-
-const vertexAI = new VertexAI({
-    project: 'adroit-crow-413218',
-    location: 'us-central1',
-    apiEndpoint: "localhost:4000/vertex-ai"
-});
-
-// Create customHeaders using Headers
-const customHeaders = new Headers({
-    "X-Litellm-Api-Key": "sk-1234"
-});
-
-// Use customHeaders in RequestOptions
-const requestOptions = {
-    customHeaders: customHeaders
-};
-
-const generativeModel = vertexAI.getGenerativeModel(
-    { model: 'gemini-1.0-pro' },
-    requestOptions
-);
-
-async function testModel() {
-    try {
-        const request = {
-            contents: [{role: 'user', parts: [{text: 'How are you doing today tell me your name?'}]}],
-          };
-        const streamingResult = await generativeModel.generateContentStream(request);
-        for await (const item of streamingResult.stream) {
-            console.log('stream chunk: ', JSON.stringify(item));
-        }
-        const aggregatedResponse = await streamingResult.response;
-        console.log('aggregated response: ', JSON.stringify(aggregatedResponse));
-    } catch (error) {
-        console.error('Error:', error);
-    }
-}
-
-testModel();
\ No newline at end of file
diff --git a/tests/pass_through_tests/test_vertex.test.js b/tests/pass_through_tests/test_vertex.test.js
new file mode 100644
index 000000000..d651650c2
--- /dev/null
+++ b/tests/pass_through_tests/test_vertex.test.js
@@ -0,0 +1,71 @@
+const { VertexAI, RequestOptions } = require('@google-cloud/vertexai');
+
+
+// Import fetch if the SDK uses it
+const originalFetch = global.fetch || require('node-fetch');
+
+// Monkey-patch the fetch used internally
+global.fetch = async function patchedFetch(url, options) {
+    // Modify the URL to use HTTP instead of HTTPS
+    if (url.startsWith('https://localhost:4000')) {
+        url = url.replace('https://', 'http://');
+    }
+    console.log('Patched fetch sending request to:', url);
+    return originalFetch(url, options);
+};
+
+
+describe('Vertex AI Tests', () => {
+    test('should successfully generate content from Vertex AI', async () => {
+        const vertexAI = new VertexAI({
+            project: 'adroit-crow-413218',
+            location: 'us-central1',
+            apiEndpoint: "localhost:4000/vertex-ai"
+        });
+
+        const customHeaders = new Headers({
+            "X-Litellm-Api-Key": "sk-1234"
+        });
+
+        const requestOptions = {
+            customHeaders: customHeaders
+        };
+
+        const generativeModel = vertexAI.getGenerativeModel(
+            { model: 'gemini-1.0-pro' },
+            requestOptions
+        );
+
+        const request = {
+            contents: [{role: 'user', parts: [{text: 'How are you doing today tell me your name?'}]}],
+        };
+
+        const streamingResult = await generativeModel.generateContentStream(request);
+        
+        // Add some assertions
+        expect(streamingResult).toBeDefined();
+        
+        for await (const item of streamingResult.stream) {
+            console.log('stream chunk:', JSON.stringify(item));
+            expect(item).toBeDefined();
+        }
+
+        const aggregatedResponse = await streamingResult.response;
+        console.log('aggregated response:', JSON.stringify(aggregatedResponse));
+        expect(aggregatedResponse).toBeDefined();
+    });
+
+
+    test('should successfully generate non-streaming content from Vertex AI', async () => {
+        const vertexAI = new VertexAI({project: 'adroit-crow-413218', location: 'us-central1', apiEndpoint: "localhost:4000/vertex-ai"});
+        const customHeaders = new Headers({"X-Litellm-Api-Key": "sk-1234"});
+        const requestOptions = {customHeaders: customHeaders};
+        const generativeModel = vertexAI.getGenerativeModel({model: 'gemini-1.0-pro'}, requestOptions);
+        const request = {contents: [{role: 'user', parts: [{text: 'What is 2+2?'}]}]};
+
+        const result = await generativeModel.generateContent(request);
+        expect(result).toBeDefined();
+        expect(result.response).toBeDefined();
+        console.log('non-streaming response:', JSON.stringify(result.response));
+    });
+});
\ No newline at end of file