batchpredict-with-bq.test.js

// 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
//
//     https://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.

'use strict';

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 = process.env.GOOGLE_CLOUD_LOCATION || 'global';
const sample = require('../batch-prediction/batchpredict-with-bq');

async function bq_output_uri() {
  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}`;

  return {
    uri: fullUri,
    async cleanup() {
      await bigquery.dataset(BQ_OUTPUT_DATASET).table(tableName).delete({
        ignoreNotFound: true,
      });
    },
  };
}

describe('batchpredict-with-bq', () => {
  it('should return the batch job state', async function () {
    this.timeout(50000);
    const bqOutput = bq_output_uri();
    const bqUri = (await bqOutput).uri;
    const output = await sample.runBatchPredictionJob(
      bqUri,
      projectId,
      location
    );
    assert.notEqual(output, undefined);
  });
});