test: add vector search unit tests

firebase · Jul 16, 2024 · 51c0c8f · 51c0c8f
1 parent afdf544
commit 51c0c8f
Show file tree

Hide file tree

Showing 5 changed files with 232 additions and 2 deletions.
diff --git a/js/plugins/vertexai/package.json b/js/plugins/vertexai/package.json
@@ -25,7 +25,7 @@
     "build:clean": "rm -rf ./lib",
     "build": "npm-run-all build:clean check compile",
     "build:watch": "tsup-node --watch",
-    "test": "tsx --test ./tests/*_test.ts"
+    "test": "tsx --test ./tests/**/*_test.ts"
   },
   "repository": {
     "type": "git",

diff --git a/js/plugins/vertexai/tests/vector-search/query_public_endpoint_test.ts b/js/plugins/vertexai/tests/vector-search/query_public_endpoint_test.ts
@@ -0,0 +1,85 @@
+/**
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import assert from 'assert';
+import test, { Mock } from 'node:test';
+import { queryPublicEndpoint } from '../../src/vector-search/query_public_endpoint'; // Replace with the actual path to your module
+
+test('queryPublicEndpoint sends the correct request and retrieves neighbors', async (t) => {
+  t.mock.method(global, 'fetch', async (url, options) => {
+    return {
+      ok: true,
+      json: async () => ({ neighbors: ['neighbor1', 'neighbor2'] }),
+    } as any;
+  });
+
+  const params = {
+    featureVector: [0.1, 0.2, 0.3],
+    neighborCount: 5,
+    accessToken: 'test-access-token',
+    projectId: 'test-project-id',
+    location: 'us-central1',
+    indexEndpointId: 'idx123',
+    publicDomainName: 'example.com',
+    projectNumber: '123456789',
+    deployedIndexId: 'deployed-idx123',
+  };
+
+  const expectedResponse = { neighbors: ['neighbor1', 'neighbor2'] };
+
+  const response = await queryPublicEndpoint(params);
+
+  type mockResponse = { neighbors: string[] };
+
+  const calls = (
+    global.fetch as Mock<
+      (url: string, options: Record<string, any>) => Promise<Response>
+    >
+  ).mock.calls;
+
+  assert.strictEqual(calls.length, 1);
+
+  const [url, options] = calls[0].arguments;
+
+  const expectedUrl = `https://example.com/v1/projects/123456789/locations/us-central1/indexEndpoints/idx123:findNeighbors`;
+
+  assert.strictEqual(url.toString(), expectedUrl);
+
+  assert.strictEqual(options.method, 'POST');
+
+  assert.strictEqual(options.headers['Content-Type'], 'application/json');
+  assert.strictEqual(
+    options.headers['Authorization'],
+    'Bearer test-access-token'
+  );
+
+  const body = JSON.parse(options.body);
+  assert.deepStrictEqual(body, {
+    deployed_index_id: 'deployed-idx123',
+    queries: [
+      {
+        datapoint: {
+          datapoint_id: '0',
+          feature_vector: [0.1, 0.2, 0.3],
+        },
+        neighbor_count: 5,
+      },
+    ],
+  });
+
+  // Verifying the response
+  assert.deepStrictEqual(response, expectedResponse);
+});
diff --git a/js/plugins/vertexai/tests/vector-search/upsert_datapoints_test.ts b/js/plugins/vertexai/tests/vector-search/upsert_datapoints_test.ts
@@ -0,0 +1,78 @@
+/**
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import assert from 'assert';
+import { GoogleAuth } from 'google-auth-library';
+import test, { Mock } from 'node:test';
+import { IIndexDatapoint } from '../../src/vector-search/types'; // Replace with the actual path to your types
+import { upsertDatapoints } from '../../src/vector-search/upsert_datapoints'; // Replace with the actual path to your module
+
+test('upsertDatapoints sends the correct request and handles response', async (t) => {
+  // Mocking the fetch method within the test scope
+  t.mock.method(global, 'fetch', async (url, options) => {
+    return {
+      ok: true,
+      json: async () => ({}),
+    } as any;
+  });
+
+  // Mocking the GoogleAuth client
+  const mockAuthClient = {
+    getAccessToken: async () => 'test-access-token',
+  } as GoogleAuth;
+
+  const params = {
+    datapoints: [
+      { datapointId: 'dp1', featureVector: [0.1, 0.2, 0.3] },
+      { datapointId: 'dp2', featureVector: [0.4, 0.5, 0.6] },
+    ] as IIndexDatapoint[],
+    authClient: mockAuthClient,
+    projectId: 'test-project-id',
+    location: 'us-central1',
+    indexId: 'idx123',
+  };
+
+  await upsertDatapoints(params);
+
+  // Verifying the fetch call
+  const calls = (
+    global.fetch as Mock<
+      (url: string, options: Record<string, any>) => Promise<Response>
+    >
+  ).mock.calls;
+
+  assert.strictEqual(calls.length, 1);
+  const [url, options] = calls[0].arguments;
+
+  assert.strictEqual(
+    url.toString(),
+    'https://us-central1-aiplatform.googleapis.com/v1/projects/test-project-id/locations/us-central1/indexes/idx123:upsertDatapoints'
+  );
+  assert.strictEqual(options.method, 'POST');
+  assert.strictEqual(options.headers['Content-Type'], 'application/json');
+  assert.strictEqual(
+    options.headers['Authorization'],
+    'Bearer test-access-token'
+  );
+
+  const body = JSON.parse(options.body);
+  assert.deepStrictEqual(body, {
+    datapoints: [
+      { datapoint_id: 'dp1', feature_vector: [0.1, 0.2, 0.3] },
+      { datapoint_id: 'dp2', feature_vector: [0.4, 0.5, 0.6] },
+    ],
+  });
+});
diff --git a/js/plugins/vertexai/tests/vector-search/utils_test.ts b/js/plugins/vertexai/tests/vector-search/utils_test.ts
@@ -0,0 +1,68 @@
+/**
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import assert from 'assert';
+import { google } from 'googleapis';
+import test from 'node:test';
+import {
+  getAccessToken,
+  getProjectNumber,
+} from '../../src/vector-search/utils'; // Replace with the actual path to your module
+
+// Mocking the google.auth.getClient method
+google.auth.getClient = async () => {
+  return {
+    getRequestHeaders: async () => ({ Authorization: 'Bearer test-token' }),
+  } as any; // Using `any` to bypass type checks for the mock
+};
+
+// Mocking the google.cloudresourcemanager method
+google.cloudresourcemanager = () => {
+  return {
+    projects: {
+      get: async ({ projectId }) => {
+        return {
+          data: {
+            projectNumber: '123456789',
+          },
+        };
+      },
+    },
+  } as any; // Using `any` to bypass type checks for the mock
+};
+
+test('getProjectNumber retrieves the project number', async () => {
+  const projectId = 'test-project-id';
+  const expectedProjectNumber = '123456789';
+
+  const projectNumber = await getProjectNumber(projectId);
+  assert.strictEqual(projectNumber, expectedProjectNumber);
+});
+
+// Mocking the GoogleAuth client
+const mockAuthClient = {
+  getAccessToken: async () => ({ token: 'test-access-token' }),
+};
+
+test('getAccessToken retrieves the access token', async () => {
+  // Mocking the GoogleAuth.getClient method to return the mockAuthClient
+  const auth = {
+    getClient: async () => mockAuthClient,
+  } as any; // Using `any` to bypass type checks for the mock
+
+  const accessToken = await getAccessToken(auth);
+  assert.strictEqual(accessToken, 'test-access-token');
+});
diff --git a/js/testapps/vertexai-vector-search-custom/documents.json b/js/testapps/vertexai-vector-search-custom/documents.json