1919import static com .google .common .truth .Truth .assertThat ;
2020import static junit .framework .TestCase .assertNotNull ;
2121
22- import com .google .cloud .dialogflow .v2beta1 .DeleteKnowledgeBaseRequest ;
23- import com .google .cloud .dialogflow .v2beta1 .KnowledgeBase ;
24- import com .google .cloud .dialogflow .v2beta1 .KnowledgeBasesClient ;
25- import com .google .cloud .dialogflow .v2beta1 . ProjectName ;
22+ import com .google .cloud .dialogflow .v2 .DeleteKnowledgeBaseRequest ;
23+ import com .google .cloud .dialogflow .v2 .KnowledgeBase ;
24+ import com .google .cloud .dialogflow .v2 .KnowledgeBasesClient ;
25+ import com .google .cloud .dialogflow .v2 . LocationName ;
2626import com .google .cloud .testing .junit4 .MultipleAttemptsRule ;
2727import java .io .ByteArrayOutputStream ;
2828import java .io .IOException ;
4141public class CreateDocumentTest {
4242
4343 private static final String PROJECT_ID = System .getenv ("GOOGLE_CLOUD_PROJECT" );
44+ private static final String LOCATION = "global" ;
4445 private static String KNOWLEDGE_DISPLAY_NAME = UUID .randomUUID ().toString ();
4546 private static String DOCUMENT_DISPLAY_NAME = UUID .randomUUID ().toString ();
46- private ByteArrayOutputStream bout ;
47- private PrintStream out ;
4847 private String knowledgeBaseName ;
48+ private ByteArrayOutputStream bout ;
49+ private PrintStream newOutputStream ;
50+ private PrintStream originalOutputStream ;
4951
5052 private static void requireEnvVar (String varName ) {
51- assertNotNull (String .format (varName ), String . format ( varName ) );
53+ assertNotNull (String .format (varName ));
5254 }
5355
5456 @ BeforeClass
@@ -59,31 +61,36 @@ public static void checkRequirements() {
5961
6062 @ Before
6163 public void setUp () throws IOException {
64+ originalOutputStream = System .out ;
65+ bout = new ByteArrayOutputStream ();
66+ newOutputStream = new PrintStream (bout );
67+ System .setOut (newOutputStream );
68+
6269 // Create a knowledge base for the document
6370 try (KnowledgeBasesClient client = KnowledgeBasesClient .create ()) {
6471 KnowledgeBase knowledgeBase =
6572 KnowledgeBase .newBuilder ().setDisplayName (KNOWLEDGE_DISPLAY_NAME ).build ();
66- ProjectName projectName = ProjectName .of (PROJECT_ID );
67- KnowledgeBase response = client .createKnowledgeBase (projectName , knowledgeBase );
73+ LocationName parent = LocationName .of (PROJECT_ID , LOCATION );
74+ KnowledgeBase response = client .createKnowledgeBase (parent , knowledgeBase );
6875 // Save the full name for deletion
6976 knowledgeBaseName = response .getName ();
7077 }
71-
72- bout = new ByteArrayOutputStream ();
73- out = new PrintStream (bout );
74- System .setOut (out );
7578 }
7679
7780 @ After
7881 public void tearDown () throws IOException {
82+ if (knowledgeBaseName == null ) {
83+ return ;
84+ }
85+
7986 // Delete the created knowledge base
8087 try (KnowledgeBasesClient client = KnowledgeBasesClient .create ()) {
8188 DeleteKnowledgeBaseRequest request =
8289 DeleteKnowledgeBaseRequest .newBuilder ().setName (knowledgeBaseName ).setForce (true ).build ();
8390 client .deleteKnowledgeBase (request );
8491 }
8592
86- System .setOut (null );
93+ System .setOut (originalOutputStream );
8794 }
8895
8996 @ Rule public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule (3 );
0 commit comments