samples: update sample code for knowledge base (#898)

* update sample code for knowledge base * chore: revert the nightly integration test project to gcloud-devel (#899) Fixes #880 Fixes #881 * sample: minor change for consistent output format * revert license year to 2018 Co-authored-by: Neenu Shaji <[email protected]>
GoogleCloudPlatform · Nov 15, 2022 · 0cd4a86 · 0cd4a86
1 parent 81a7967
commit 0cd4a86
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 26 deletions.
diff --git a/dialogflow/snippets/src/main/java/com/example/dialogflow/ConversationProfileManagement.java b/dialogflow/snippets/src/main/java/com/example/dialogflow/ConversationProfileManagement.java
@@ -125,7 +125,7 @@ public static void createConversationProfileArticleSuggestion(
                   .setConversationProfile(targetConversationProfile)
                   .build());
       System.out.println("====================");
-      System.out.println("Conversation Profile created:");
+      System.out.println("Conversation Profile created:\n");
       System.out.format("Display name: %s\n", createdConversationProfile.getDisplayName());
       System.out.format("Name: %s\n", createdConversationProfile.getName());
     }

diff --git a/dialogflow/snippets/src/main/java/com/example/dialogflow/KnowledgeBaseManagement.java b/dialogflow/snippets/src/main/java/com/example/dialogflow/KnowledgeBaseManagement.java
@@ -19,26 +19,39 @@
 // [START dialogflow_create_knowledge_base]
 
 import com.google.api.gax.rpc.ApiException;
-import com.google.cloud.dialogflow.v2beta1.KnowledgeBase;
-import com.google.cloud.dialogflow.v2beta1.KnowledgeBasesClient;
-import com.google.cloud.dialogflow.v2beta1.ProjectName;
+import com.google.cloud.dialogflow.v2.KnowledgeBase;
+import com.google.cloud.dialogflow.v2.KnowledgeBasesClient;
+import com.google.cloud.dialogflow.v2.LocationName;
 import java.io.IOException;
 
 public class KnowledgeBaseManagement {
 
+  public static void main(String[] args) throws ApiException, IOException {
+    // TODO(developer): Replace these variables before running the sample.
+    String projectId = "my-project-id";
+    String location = "my-location";
+
+    // Set display name of the new knowledge base
+    String knowledgeBaseDisplayName = "my-knowledge-base-display-name";
+
+    // Create a knowledge base
+    createKnowledgeBase(projectId, location, knowledgeBaseDisplayName);
+  }
+
   // Create a Knowledge base
-  public static KnowledgeBase createKnowledgeBase(String projectId, String displayName)
+  public static void createKnowledgeBase(String projectId, String location, String displayName)
       throws ApiException, IOException {
     // Instantiates a client
     try (KnowledgeBasesClient knowledgeBasesClient = KnowledgeBasesClient.create()) {
-      KnowledgeBase knowledgeBase = KnowledgeBase.newBuilder().setDisplayName(displayName).build();
-      ProjectName projectName = ProjectName.of(projectId);
-      KnowledgeBase response = knowledgeBasesClient.createKnowledgeBase(projectName, knowledgeBase);
+      KnowledgeBase targetKnowledgeBase =
+          KnowledgeBase.newBuilder().setDisplayName(displayName).build();
+      LocationName parent = LocationName.of(projectId, location);
+      KnowledgeBase createdKnowledgeBase =
+          knowledgeBasesClient.createKnowledgeBase(parent, targetKnowledgeBase);
+      System.out.println("====================");
       System.out.format("Knowledgebase created:\n");
-      System.out.format("Display Name: %s \n", response.getDisplayName());
-      System.out.format("Knowledge ID: %s \n", response.getName());
-
-      return response;
+      System.out.format("Display Name: %s\n", createdKnowledgeBase.getDisplayName());
+      System.out.format("Name: %s\n", createdKnowledgeBase.getName());
     }
   }
 }

diff --git a/dialogflow/snippets/src/test/java/com/example/dialogflow/CreateKnowledgeBaseTest.java b/dialogflow/snippets/src/test/java/com/example/dialogflow/CreateKnowledgeBaseTest.java
@@ -19,9 +19,8 @@
 import static com.google.common.truth.Truth.assertThat;
 import static junit.framework.TestCase.assertNotNull;
 
-import com.google.cloud.dialogflow.v2beta1.DeleteKnowledgeBaseRequest;
-import com.google.cloud.dialogflow.v2beta1.KnowledgeBase;
-import com.google.cloud.dialogflow.v2beta1.KnowledgeBasesClient;
+import com.google.cloud.dialogflow.v2.DeleteKnowledgeBaseRequest;
+import com.google.cloud.dialogflow.v2.KnowledgeBasesClient;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
@@ -38,13 +37,23 @@
 public class CreateKnowledgeBaseTest {
 
   private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT");
+  private static final String LOCATION = "global";
+  private static final String ID_PREFIX_IN_OUTPUT = "Name: ";
   private static String KNOWLEDGE_DISPLAY_NAME = UUID.randomUUID().toString();
-  private ByteArrayOutputStream bout;
-  private PrintStream out;
   private String knowledgeBaseName;
+  private ByteArrayOutputStream bout;
+  private PrintStream newOutputStream;
+  private PrintStream originalOutputStream;
 
   private static void requireEnvVar(String varName) {
-    assertNotNull(String.format(varName), String.format(varName));
+    assertNotNull(System.getenv(varName));
+  }
+
+  // Extract the name of created resource from "Name: %s\n" in sample code output
+  private static String getResourceNameFromOutputString(String output) {
+    return output.substring(
+        output.lastIndexOf(ID_PREFIX_IN_OUTPUT) + ID_PREFIX_IN_OUTPUT.length(),
+        output.length() - 1);
   }
 
   @BeforeClass
@@ -56,27 +65,30 @@ public static void checkRequirements() {
   @Before
   public void setUp() {
     bout = new ByteArrayOutputStream();
-    out = new PrintStream(bout);
-    System.setOut(out);
+    newOutputStream = new PrintStream(bout);
+    System.setOut(newOutputStream);
   }
 
   @After
   public void tearDown() throws IOException {
+    if (knowledgeBaseName == null) {
+      return;
+    }
+
     // Delete the created knowledge base
     try (KnowledgeBasesClient client = KnowledgeBasesClient.create()) {
       DeleteKnowledgeBaseRequest request =
           DeleteKnowledgeBaseRequest.newBuilder().setName(knowledgeBaseName).setForce(true).build();
       client.deleteKnowledgeBase(request);
     }
-    System.setOut(null);
+    System.setOut(originalOutputStream);
   }
 
   @Test
   public void testCreateKnowledgeBase() throws Exception {
-    KnowledgeBase knowledgeBase =
-        KnowledgeBaseManagement.createKnowledgeBase(PROJECT_ID, KNOWLEDGE_DISPLAY_NAME);
-    knowledgeBaseName = knowledgeBase.getName();
-    String got = bout.toString();
-    assertThat(got).contains(KNOWLEDGE_DISPLAY_NAME);
+    KnowledgeBaseManagement.createKnowledgeBase(PROJECT_ID, LOCATION, KNOWLEDGE_DISPLAY_NAME);
+    String output = bout.toString();
+    assertThat(output).contains(KNOWLEDGE_DISPLAY_NAME);
+    knowledgeBaseName = getResourceNameFromOutputString(output);
   }
 }