Merge pull request #3290 from GoogleCloudPlatform/list-llm-test-fix

fix(aiplatform): list tuned models test
GoogleCloudPlatform · Jun 22, 2023 · 62b70e8 · 62b70e8
2 parents 7a360b6 + 6fe10d2
commit 62b70e8
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/ai-platform/snippets/list-tuned-models.js b/ai-platform/snippets/list-tuned-models.js
@@ -15,7 +15,7 @@
  */
 'use strict';
 
-function main(project, location, model = 'text-bison@001') {
+async function main(project, location, model = 'text-bison-001') {
   // [START aiplatform_list_tuned_models]
   /**
    * TODO(developer): Uncomment these variables before running the sample.\
@@ -34,7 +34,7 @@ function main(project, location, model = 'text-bison@001') {
   // Instantiate the service client.
   const modelServiceClient = new ModelServiceClient(clientOptions);
 
-  function listTunedModels() {
+  async function listTunedModels() {
     // Configure the parent resource
     const parent = `projects/${project}/locations/${location}`;
     const filter = `labels.google-vertex-llm-tuning-base-model-id=${model}`;
@@ -44,14 +44,14 @@ function main(project, location, model = 'text-bison@001') {
       filter,
     };
 
-    const [response] = modelServiceClient.listModels(request);
+    const [response] = await modelServiceClient.listModels(request);
     console.log('List Tuned Models response');
     for (const model of response) {
       console.log(`\tModel name: ${model.name}`);
       console.log(`\tDisplay name: ${model.displayName}`);
     }
   }
-  listTunedModels();
+  await listTunedModels();
   // [END aiplatform_list_tuned_models]
 }
 

diff --git a/ai-platform/snippets/test/list-tuned-models.test.js b/ai-platform/snippets/test/list-tuned-models.test.js
@@ -25,7 +25,7 @@ const project = process.env.CAIP_PROJECT_ID;
 
 const {listTunedModels} = require('../list-tuned-models');
 
-describe('AI platform list tuned models', () => {
+describe('List tuned models', async () => {
   const stubConsole = function () {
     sinon.stub(console, 'error');
     sinon.stub(console, 'log');
@@ -39,8 +39,8 @@ describe('AI platform list tuned models', () => {
   beforeEach(stubConsole);
   afterEach(restoreConsole);
 
-  it('should list all tuned LLM models', () => {
-    listTunedModels(project, LOCATION);
+  it('should list all tuned LLM models', async () => {
+    await listTunedModels(project, LOCATION);
     assert.include(console.log.firstCall.args, 'List Tuned Models response');
   });
 });