1919import static com .google .common .truth .Truth .assertThat ;
2020import static junit .framework .TestCase .assertNotNull ;
2121
22+ import com .google .cloud .testing .junit4 .MultipleAttemptsRule ;
2223import java .io .ByteArrayOutputStream ;
2324import java .io .IOException ;
2425import java .io .PrintStream ;
2728import org .junit .After ;
2829import org .junit .Before ;
2930import org .junit .BeforeClass ;
31+ import org .junit .Rule ;
3032import org .junit .Test ;
3133import org .junit .runner .RunWith ;
3234import org .junit .runners .JUnit4 ;
@@ -39,7 +41,7 @@ public class TranslateCreateDatasetTest {
3941 private ByteArrayOutputStream bout ;
4042 private PrintStream out ;
4143 private PrintStream originalPrintStream ;
42- private String datasetId ;
44+ private String got ;
4345
4446 private static void requireEnvVar (String varName ) {
4547 assertNotNull (
@@ -63,11 +65,15 @@ public void setUp() {
6365
6466 @ After
6567 public void tearDown () throws InterruptedException , ExecutionException , IOException {
68+ String datasetId = got .split ("Dataset id: " )[1 ].split ("\n " )[0 ];
69+
6670 // Delete the created dataset
6771 DeleteDataset .deleteDataset (PROJECT_ID , datasetId );
6872 System .setOut (originalPrintStream );
6973 }
7074
75+ @ Rule public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule (3 );
76+
7177 @ Test
7278 public void testCreateDataset () throws IOException , ExecutionException , InterruptedException {
7379 // Create a random dataset name with a length of 32 characters (max allowed by AutoML)
@@ -77,8 +83,7 @@ public void testCreateDataset() throws IOException, ExecutionException, Interrup
7783 String .format ("test_%s" , UUID .randomUUID ().toString ().replace ("-" , "_" ).substring (0 , 26 ));
7884 TranslateCreateDataset .createDataset (PROJECT_ID , datasetName );
7985
80- String got = bout .toString ();
86+ got = bout .toString ();
8187 assertThat (got ).contains ("Dataset id:" );
82- datasetId = got .split ("Dataset id: " )[1 ].split ("\n " )[0 ];
8388 }
8489}
0 commit comments