test(ai): add integration tests

.gitignore

@@ -103,4 +103,7 @@ vertexai-sdk-test-data
 mocks-lookup.ts
 
 # temp changeset output
-changeset-temp.json
+changeset-temp.json
+
+# Vertex AI integration test Firebase project config
+packages/vertexai/integration/firebase-config.ts

packages/vertexai/integration/constants.ts

@@ -0,0 +1,63 @@
+/**
+ * @license
+ * Copyright 2025 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 {
+  Content,
+  GenerationConfig,
+  HarmBlockMethod,
+  HarmBlockThreshold,
+  HarmCategory,
+  SafetySetting
+} from '../src';
+
+export const MODEL_NAME = 'gemini-1.5-pro';
+
+export const generationConfig: GenerationConfig = {
+  temperature: 0,
+  topP: 0,
+  responseMimeType: 'text/plain'
+};
+
+export const safetySettings: SafetySetting[] = [
+  {
+    category: HarmCategory.HARM_CATEGORY_HARASSMENT,
+    threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+    method: HarmBlockMethod.PROBABILITY
+  },
+  {
+    category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+    threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+    method: HarmBlockMethod.SEVERITY
+  },
+  {
+    category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+    threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
+  },
+  {
+    category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+    threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
+  }
+];
+
+export const systemInstruction: Content = {
+  role: 'system',
+  parts: [
+    {
+      text: 'You are a friendly and helpful assistant.'
+    }
+  ]
+};

packages/vertexai/integration/count-tokens.test.ts

@@ -0,0 +1,56 @@
+/**
+ * @license
+ * Copyright 2025 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 { expect } from 'chai';
+import { Modality, getGenerativeModel, getVertexAI } from '../src';
+import {
+  MODEL_NAME,
+  generationConfig,
+  systemInstruction,
+  safetySettings
+} from './constants';
+import { initializeApp } from '@firebase/app';
+import { FIREBASE_CONFIG } from './firebase-config';
+
+describe('Count Tokens', () => {
+  before(() => initializeApp(FIREBASE_CONFIG));
+
+  it('CountTokens text', async () => {
+    const vertexAI = getVertexAI();
+    const model = getGenerativeModel(vertexAI, {
+      model: MODEL_NAME,
+      generationConfig,
+      systemInstruction,
+      safetySettings
+    });
+
+    const response = await model.countTokens('Why is the sky blue?');
+
+    expect(response.totalTokens).to.equal(6);
+    expect(response.totalBillableCharacters).to.equal(16);
+    expect(response.promptTokensDetails).to.not.be.null;
+    expect(response.promptTokensDetails!.length).to.equal(1);
+    expect(response.promptTokensDetails![0].modality).to.equal(Modality.TEXT);
+    expect(response.promptTokensDetails![0].tokenCount).to.equal(6);
+  });
+  // TODO (dlarocque): Test countTokens() with the following:
+  // - inline data
+  // - public storage reference
+  // - private storage reference (testing auth integration)
+  // - count tokens
+  // - JSON schema
+});

packages/vertexai/integration/firebase-config.ts

@@ -0,0 +1,19 @@
+/**
+ * @license
+ * Copyright 2025 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.
+ */
+
+// PLACEHOLDER: This should be replaced with a Firebase config for a project with access to Vertex AI.
+export const FIREBASE_CONFIG = {};

packages/vertexai/integration/generate-content.test.ts

@@ -0,0 +1,84 @@
+/**
+ * @license
+ * Copyright 2025 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 { expect } from 'chai';
+import { Modality, getGenerativeModel, getVertexAI } from '../src';
+import {
+  MODEL_NAME,
+  generationConfig,
+  systemInstruction,
+  safetySettings
+} from './constants';
+import { initializeApp } from '@firebase/app';
+import { FIREBASE_CONFIG } from './firebase-config';
+
+// Token counts are only expected to differ by at most this number of tokens.
+// Set to 1 for whitespace that is not always present.
+const TOKEN_COUNT_DELTA = 1;
+
+describe('Generate Content', () => {
+  before(() => initializeApp(FIREBASE_CONFIG));
+
+  it('generateContent', async () => {
+    const vertexAI = getVertexAI();
+    const model = getGenerativeModel(vertexAI, {
+      model: MODEL_NAME,
+      generationConfig,
+      systemInstruction,
+      safetySettings
+    });
+
+    const result = await model.generateContent(
+      'Where is Google headquarters located? Answer with the city name only.'
+    );
+    const response = result.response;
+
+    const trimmedText = response.text().trim();
+    expect(trimmedText).to.equal('Mountain View');
+
+    expect(response.usageMetadata).to.not.be.null;
+    expect(response.usageMetadata!.promptTokenCount).to.be.closeTo(
+      21,
+      TOKEN_COUNT_DELTA
+    );
+    expect(response.usageMetadata!.candidatesTokenCount).to.be.closeTo(
+      4,
+      TOKEN_COUNT_DELTA
+    );
+    expect(response.usageMetadata!.totalTokenCount).to.be.closeTo(
+      25,
+      TOKEN_COUNT_DELTA * 2
+    );
+    expect(response.usageMetadata!.promptTokensDetails).to.not.be.null;
+    expect(response.usageMetadata!.promptTokensDetails!.length).to.equal(1);
+    expect(response.usageMetadata!.promptTokensDetails![0].modality).to.equal(
+      Modality.TEXT
+    );
+    expect(response.usageMetadata!.promptTokensDetails![0].tokenCount).to.equal(
+      21
+    );
+    expect(response.usageMetadata!.candidatesTokensDetails).to.not.be.null;
+    expect(response.usageMetadata!.candidatesTokensDetails!.length).to.equal(1);
+    expect(
+      response.usageMetadata!.candidatesTokensDetails![0].modality
+    ).to.equal(Modality.TEXT);
+    expect(
+      response.usageMetadata!.candidatesTokensDetails![0].tokenCount
+    ).to.equal(4);
+  });
+  // TODO (dlarocque): Test generateContentStream
+});

packages/vertexai/karma.conf.js

@@ -16,14 +16,38 @@
  */
 
 const karmaBase = require('../../config/karma.base');
+const { argv } = require('yargs');
+const { existsSync } = require('fs');
 
 const files = [`src/**/*.test.ts`];
 
+// Validate that the file that defines the Firebase config to be used in the integration tests exists.
+if (argv.integration) {
+  if (!existsSync('integration/firebase-config.ts')) {
+    throw new Error(
+      `integration/firebase-config.ts does not exist. This file must contain a Firebase config for a project with Vertex AI enabled.`
+    );
+  }
+}
+
 module.exports = function (config) {
   const karmaConfig = {
     ...karmaBase,
+
+    preprocessors: {
+      ...karmaBase.preprocessors,
+      'integration/**/*.ts': ['webpack', 'sourcemap']
+    },
+
     // files to load into karma
-    files,
+    files: (() => {
+      if (argv.integration) {
+        return ['integration/**'];
+      } else {
+        return ['src/**/*.test.ts'];
+      }
+    })(),
+
     // frameworks to use
     // available frameworks: https://npmjs.org/browse/keyword/karma-adapter
     frameworks: ['mocha']

packages/vertexai/package.json

@@ -39,6 +39,7 @@
     "test:ci": "yarn testsetup && node ../../scripts/run_tests_in_ci.js -s test",
     "test:skip-clone": "karma start",
     "test:browser": "yarn testsetup && karma start",
+    "test:integration": "karma start --integration",
     "api-report": "api-extractor run --local --verbose",
     "typings:public": "node ../../scripts/build/use_typings.js ./dist/vertexai-public.d.ts",
     "trusted-type-check": "tsec -p tsconfig.json --noEmit"