From 255b3852623e37c62b5f97bc486ba9a342b86c5b Mon Sep 17 00:00:00 2001
From: Eric Schmidt <erschmid@google.com>
Date: Thu, 25 Mar 2021 09:24:04 -0700
Subject: [PATCH] fix(samples): adds retry to TranslateCreateDatasetTest (#513)

* fix(samples): adds retry to TranslateCreateDatasetTest
---
 automl/snippets/pom.xml                               |  7 +++++++
 .../example/automl/TranslateCreateDatasetTest.java    | 11 ++++++++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/automl/snippets/pom.xml b/automl/snippets/pom.xml
index 1bd38367e5d..d87ec1b48f4 100644
--- a/automl/snippets/pom.xml
+++ b/automl/snippets/pom.xml
@@ -68,6 +68,13 @@
       <version>1.1.2</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>google-cloud-core</artifactId>
+      <version>1.94.0</version>
+      <scope>test</scope>
+      <classifier>tests</classifier>
+    </dependency>
     <!-- [END_EXCLUDE] -->
   </dependencies>
   <!-- [END automl_install_with_bom] -->
diff --git a/automl/snippets/src/test/java/com/example/automl/TranslateCreateDatasetTest.java b/automl/snippets/src/test/java/com/example/automl/TranslateCreateDatasetTest.java
index ad2081a06ad..6fd75eff945 100644
--- a/automl/snippets/src/test/java/com/example/automl/TranslateCreateDatasetTest.java
+++ b/automl/snippets/src/test/java/com/example/automl/TranslateCreateDatasetTest.java
@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static junit.framework.TestCase.assertNotNull;
 
+import com.google.cloud.testing.junit4.MultipleAttemptsRule;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
@@ -27,6 +28,7 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -39,7 +41,7 @@ public class TranslateCreateDatasetTest {
   private ByteArrayOutputStream bout;
   private PrintStream out;
   private PrintStream originalPrintStream;
-  private String datasetId;
+  private String got;
 
   private static void requireEnvVar(String varName) {
     assertNotNull(
@@ -63,11 +65,15 @@ public void setUp() {
 
   @After
   public void tearDown() throws InterruptedException, ExecutionException, IOException {
+    String datasetId = got.split("Dataset id: ")[1].split("\n")[0];
+
     // Delete the created dataset
     DeleteDataset.deleteDataset(PROJECT_ID, datasetId);
     System.setOut(originalPrintStream);
   }
 
+  @Rule public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule(3);
+
   @Test
   public void testCreateDataset() throws IOException, ExecutionException, InterruptedException {
     // 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
         String.format("test_%s", UUID.randomUUID().toString().replace("-", "_").substring(0, 26));
     TranslateCreateDataset.createDataset(PROJECT_ID, datasetName);
 
-    String got = bout.toString();
+    got = bout.toString();
     assertThat(got).contains("Dataset id:");
-    datasetId = got.split("Dataset id: ")[1].split("\n")[0];
   }
 }