diff --git a/dialogflow/snippets/src/main/java/com/example/dialogflow/DocumentManagement.java b/dialogflow/snippets/src/main/java/com/example/dialogflow/DocumentManagement.java
index a78e7ae058b..08431689008 100644
--- a/dialogflow/snippets/src/main/java/com/example/dialogflow/DocumentManagement.java
+++ b/dialogflow/snippets/src/main/java/com/example/dialogflow/DocumentManagement.java
@@ -20,11 +20,11 @@
 
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.api.gax.rpc.ApiException;
-import com.google.cloud.dialogflow.v2beta1.CreateDocumentRequest;
-import com.google.cloud.dialogflow.v2beta1.Document;
-import com.google.cloud.dialogflow.v2beta1.Document.KnowledgeType;
-import com.google.cloud.dialogflow.v2beta1.DocumentsClient;
-import com.google.cloud.dialogflow.v2beta1.KnowledgeOperationMetadata;
+import com.google.cloud.dialogflow.v2.CreateDocumentRequest;
+import com.google.cloud.dialogflow.v2.Document;
+import com.google.cloud.dialogflow.v2.Document.KnowledgeType;
+import com.google.cloud.dialogflow.v2.DocumentsClient;
+import com.google.cloud.dialogflow.v2.KnowledgeOperationMetadata;
 import java.io.IOException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -32,7 +32,7 @@
 
 public class DocumentManagement {
 
-  public static Document createDocument(
+  public static void createDocument(
       String knowledgeBaseName,
       String displayName,
       String mimeType,
@@ -58,14 +58,13 @@ public static Document createDocument(
       Document createdDocument = response.get(180, TimeUnit.SECONDS);
       System.out.format("Created Document:\n");
       System.out.format(" - Display Name: %s\n", createdDocument.getDisplayName());
-      System.out.format(" - Knowledge ID: %s\n", createdDocument.getName());
+      System.out.format(" - Document Name: %s\n", createdDocument.getName());
       System.out.format(" - MIME Type: %s\n", createdDocument.getMimeType());
       System.out.format(" - Knowledge Types:\n");
       for (KnowledgeType knowledgeTypeId : document.getKnowledgeTypesList()) {
         System.out.format("  - %s \n", knowledgeTypeId.getValueDescriptor());
       }
       System.out.format(" - Source: %s \n", document.getContentUri());
-      return createdDocument;
     }
   }
 }
diff --git a/dialogflow/snippets/src/test/java/com/example/dialogflow/CreateDocumentTest.java b/dialogflow/snippets/src/test/java/com/example/dialogflow/CreateDocumentTest.java
index 41c8d17e8b5..e59508a521e 100644
--- a/dialogflow/snippets/src/test/java/com/example/dialogflow/CreateDocumentTest.java
+++ b/dialogflow/snippets/src/test/java/com/example/dialogflow/CreateDocumentTest.java
@@ -19,10 +19,10 @@
 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.v2beta1.ProjectName;
+import com.google.cloud.dialogflow.v2.DeleteKnowledgeBaseRequest;
+import com.google.cloud.dialogflow.v2.KnowledgeBase;
+import com.google.cloud.dialogflow.v2.KnowledgeBasesClient;
+import com.google.cloud.dialogflow.v2.LocationName;
 import com.google.cloud.testing.junit4.MultipleAttemptsRule;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -41,14 +41,16 @@
 public class CreateDocumentTest {
 
   private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT");
+  private static final String LOCATION = "global";
   private static String KNOWLEDGE_DISPLAY_NAME = UUID.randomUUID().toString();
   private static String DOCUMENT_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(String.format(varName));
   }
 
   @BeforeClass
@@ -59,23 +61,28 @@ public static void checkRequirements() {
 
   @Before
   public void setUp() throws IOException {
+    originalOutputStream = System.out;
+    bout = new ByteArrayOutputStream();
+    newOutputStream = new PrintStream(bout);
+    System.setOut(newOutputStream);
+
     // Create a knowledge base for the document
     try (KnowledgeBasesClient client = KnowledgeBasesClient.create()) {
       KnowledgeBase knowledgeBase =
           KnowledgeBase.newBuilder().setDisplayName(KNOWLEDGE_DISPLAY_NAME).build();
-      ProjectName projectName = ProjectName.of(PROJECT_ID);
-      KnowledgeBase response = client.createKnowledgeBase(projectName, knowledgeBase);
+      LocationName parent = LocationName.of(PROJECT_ID, LOCATION);
+      KnowledgeBase response = client.createKnowledgeBase(parent, knowledgeBase);
       // Save the full name for deletion
       knowledgeBaseName = response.getName();
     }
-
-    bout = new ByteArrayOutputStream();
-    out = new PrintStream(bout);
-    System.setOut(out);
   }
 
   @After
   public void tearDown() throws IOException {
+    if (knowledgeBaseName == null) {
+      return;
+    }
+
     // Delete the created knowledge base
     try (KnowledgeBasesClient client = KnowledgeBasesClient.create()) {
       DeleteKnowledgeBaseRequest request =
@@ -83,7 +90,7 @@ public void tearDown() throws IOException {
       client.deleteKnowledgeBase(request);
     }
 
-    System.setOut(null);
+    System.setOut(originalOutputStream);
   }
 
   @Rule public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule(3);