adding mocking for batch prediction

Author: Guiners

genai/batch-prediction/batchpredict-embeddings-with-gcs.js

@@ -64,10 +64,6 @@ async function runBatchPredictionJob(
     await new Promise(resolve => setTimeout(resolve, 30000));
     job = await client.batches.get({name: job.name});
     console.log(`Job state: ${job.state}`);
-    // if (job.state === 'JOB_STATE_FAILED') {
-    //   console.log(`Error: ${job.state}`);
-    //   break;
-    // }
   }
 
   // Example response:

genai/test/batchpredict-embeddings-with-gcs.test.js

@@ -16,40 +16,34 @@
 
 const {assert} = require('chai');
 const {describe, it} = require('mocha');
-const {Storage} = require('@google-cloud/storage');
-
-const storage = new Storage();
-
-const GCS_OUTPUT_BUCKET = 'nodejs-docs-samples-tests';
 
 const projectId = process.env.CAIP_PROJECT_ID;
 const location = 'us-central1';
-const sample = require('../batch-prediction/batchpredict-embeddings-with-gcs');
 const {delay} = require('./util');
+const proxyquire = require('proxyquire');
+const {GoogleGenAI_Mock} = require('./batchprediction-utils');
+
+const sample = proxyquire(
+  '../batch-prediction/batchpredict-embeddings-with-gcs',
+  {
+    '@google/genai': {
+      GoogleGenAI: GoogleGenAI_Mock,
+    },
+  }
+);
 
 async function getGcsOutputUri() {
-  const dt = new Date();
-  const prefix = `text_output/${dt.toISOString()}`;
-  const fullUri = `gs://${GCS_OUTPUT_BUCKET}/${prefix}`;
-
   return {
-    uri: fullUri,
-    async cleanup() {
-      const [files] = await storage.bucket(GCS_OUTPUT_BUCKET).getFiles({
-        prefix,
-      });
-      for (const file of files) {
-        await file.delete();
-      }
-    },
+    uri: 'gs://mock/output',
+    async cleanup() {},
   };
 }
 
 describe('batchpredict-with-gcs', () => {
   it('should return the batch job state', async function () {
     this.timeout(500000);
-    // this.retries(4);
-    // await delay(this.test);
+    this.retries(4);
+    await delay(this.test);
     const gcsOutput = await getGcsOutputUri();
     try {
       const output = await sample.runBatchPredictionJob(

genai/test/batchpredict-with-bq.test.js

@@ -16,31 +16,24 @@
 
 const {assert} = require('chai');
 const {describe, it} = require('mocha');
-const {BigQuery} = require('@google-cloud/bigquery');
-
-const bigquery = new BigQuery();
-
-const BQ_OUTPUT_DATASET = `${process.env.BQ_OUTPUT_DATASET}.gen_ai_batch_prediction`;
 
 const projectId = process.env.CAIP_PROJECT_ID;
 const location = 'us-central1';
 const {delay} = require('./util');
 
-const sample = require('../batch-prediction/batchpredict-with-bq');
+const proxyquire = require('proxyquire');
+const {GoogleGenAI_Mock} = require('./batchprediction-utils');
 
-async function getBqOutputUri() {
-  const dt = new Date();
-  const tableName = `text_output_${dt.getFullYear()}_${dt.getMonth() + 1}_${dt.getDate()}_T${dt.getHours()}_${dt.getMinutes()}_${dt.getSeconds()}`;
-  const tableUri = `${BQ_OUTPUT_DATASET}.${tableName}`;
-  const fullUri = `bq://${tableUri}`;
+const sample = proxyquire('../batch-prediction/batchpredict-with-bq', {
+  '@google/genai': {
+    GoogleGenAI: GoogleGenAI_Mock,
+  },
+});
 
+async function getBqOutputUri() {
   return {
-    uri: fullUri,
-    async cleanup() {
-      await bigquery.dataset(BQ_OUTPUT_DATASET).table(tableName).delete({
-        ignoreNotFound: true,
-      });
-    },
+    uri: 'gs://mock/output',
+    async cleanup() {},
   };
 }
 

genai/test/batchpredict-with-gcs.test.js

@@ -16,48 +16,43 @@
 
 const {assert} = require('chai');
 const {describe, it} = require('mocha');
-const {Storage} = require('@google-cloud/storage');
-
-const storage = new Storage();
-
-const GCS_OUTPUT_BUCKET = 'nodejs-docs-samples-tests';
+const proxyquire = require('proxyquire');
 
 const projectId = process.env.CAIP_PROJECT_ID;
 const location = 'us-central1';
-const sample = require('../batch-prediction/batchpredict-with-gcs');
 const {delay} = require('./util');
+const {GoogleGenAI_Mock} = require('./batchprediction-utils');
 
-async function getGcsOutputUri() {
-  const dt = new Date();
-  const prefix = `text_output/${dt.toISOString()}`;
-  const fullUri = `gs://${GCS_OUTPUT_BUCKET}/${prefix}`;
+const sample = proxyquire('../batch-prediction/batchpredict-with-gcs', {
+  '@google/genai': {
+    GoogleGenAI: GoogleGenAI_Mock,
+  },
+});
 
+async function getGcsOutputUri() {
   return {
-    uri: fullUri,
-    async cleanup() {
-      const [files] = await storage.bucket(GCS_OUTPUT_BUCKET).getFiles({
-        prefix,
-      });
-      for (const file of files) {
-        await file.delete();
-      }
-    },
+    uri: 'gs://mock/output',
+    async cleanup() {},
   };
 }
 
-describe('batchpredict-with-gcs', () => {
+describe('batchpredict-with-gcs (mocked)', () => {
   it('should return the batch job state', async function () {
     this.timeout(500000);
     this.retries(4);
     await delay(this.test);
+
     const gcsOutput = await getGcsOutputUri();
+
     try {
       const output = await sample.runBatchPredictionJob(
         gcsOutput.uri,
         projectId,
         location
       );
-      assert.notEqual(output, undefined);
+
+      console.log('output', output);
+      assert.equal(output, 'JOB_STATE_SUCCEEDED');
     } finally {
       await gcsOutput.cleanup();
     }

genai/test/batchprediction-utils.js

@@ -0,0 +1,32 @@
+const GoogleGenAI_Mock = function () {
+  let getCalled = false;
+
+  return {
+    batches: {
+      create: async () => ({
+        name: 'projects/mock/locations/mock/batchPredictionJobs/123',
+        state: 'JOB_STATE_PENDING',
+      }),
+
+      get: async () => {
+        // First call returns running, second call returns success
+        if (!getCalled) {
+          getCalled = true;
+          return {
+            name: 'projects/mock/locations/mock/batchPredictionJobs/123',
+            state: 'JOB_STATE_RUNNING',
+          };
+        }
+
+        return {
+          name: 'projects/mock/locations/mock/batchPredictionJobs/123',
+          state: 'JOB_STATE_SUCCEEDED',
+        };
+      },
+    },
+  };
+};
+
+module.exports = {
+  GoogleGenAI_Mock,
+};