test: Increased the featurestore tests timeout (GoogleCloudPlatform#7700)

gericdong · web-flow · commit 585fac41307b · 2023-02-22T16:48:30.000-05:00
* test: increased the featurestore tests timeout

* enabled the featurestore tests

* Explicitly set the LRO timeout

* Lint reformat
diff --git a/aiplatform/src/main/java/aiplatform/CreateFeaturestoreFixedNodesSample.java b/aiplatform/src/main/java/aiplatform/CreateFeaturestoreFixedNodesSample.java
@@ -24,17 +24,21 @@
 // [START aiplatform_create_featurestore_fixed_nodes_sample]
 
 import com.google.api.gax.longrunning.OperationFuture;
+import com.google.api.gax.longrunning.OperationTimedPollAlgorithm;
+import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.aiplatform.v1beta1.CreateFeaturestoreOperationMetadata;
 import com.google.cloud.aiplatform.v1beta1.CreateFeaturestoreRequest;
 import com.google.cloud.aiplatform.v1beta1.Featurestore;
 import com.google.cloud.aiplatform.v1beta1.Featurestore.OnlineServingConfig;
 import com.google.cloud.aiplatform.v1beta1.FeaturestoreServiceClient;
 import com.google.cloud.aiplatform.v1beta1.FeaturestoreServiceSettings;
 import com.google.cloud.aiplatform.v1beta1.LocationName;
+import com.google.cloud.aiplatform.v1beta1.stub.FeaturestoreServiceStubSettings;
 import java.io.IOException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import org.threeten.bp.Duration;
 
 public class CreateFeaturestoreFixedNodesSample {
 
@@ -60,8 +64,30 @@ static void createFeaturestoreFixedNodesSample(
       int timeout)
       throws IOException, InterruptedException, ExecutionException, TimeoutException {
 
+    OperationTimedPollAlgorithm operationTimedPollAlgorithm =
+        OperationTimedPollAlgorithm.create(
+            RetrySettings.newBuilder()
+                .setInitialRetryDelay(Duration.ofMillis(5000L))
+                .setRetryDelayMultiplier(1.5)
+                .setMaxRetryDelay(Duration.ofMillis(45000L))
+                .setInitialRpcTimeout(Duration.ZERO)
+                .setRpcTimeoutMultiplier(1.0)
+                .setMaxRpcTimeout(Duration.ZERO)
+                .setTotalTimeout(Duration.ofSeconds(timeout))
+                .build());
+
+    FeaturestoreServiceStubSettings.Builder featurestoreServiceStubSettingsBuilder =
+        FeaturestoreServiceStubSettings.newBuilder();
+
+    featurestoreServiceStubSettingsBuilder
+        .createFeaturestoreOperationSettings()
+        .setPollingAlgorithm(operationTimedPollAlgorithm);
+    FeaturestoreServiceStubSettings featureStoreStubSettings =
+        featurestoreServiceStubSettingsBuilder.build();
     FeaturestoreServiceSettings featurestoreServiceSettings =
-        FeaturestoreServiceSettings.newBuilder().setEndpoint(endpoint).build();
+        FeaturestoreServiceSettings.create(featureStoreStubSettings);
+    featurestoreServiceSettings =
+        featurestoreServiceSettings.toBuilder().setEndpoint(endpoint).build();
 
     // Initialize client that will be used to send requests. This client only needs to be created
     // once, and can be reused for multiple requests. After completing all of your requests, call
diff --git a/aiplatform/src/test/java/aiplatform/FeaturestoreSamplesTest.java b/aiplatform/src/test/java/aiplatform/FeaturestoreSamplesTest.java
@@ -49,7 +49,7 @@ public class FeaturestoreSamplesTest {
   private static final boolean USE_FORCE = true;
   private static final String LOCATION = "us-central1";
   private static final String ENDPOINT = "us-central1-aiplatform.googleapis.com:443";
-  private static final int TIMEOUT = 900;
+  private static final int TIMEOUT = 1800;
   private ByteArrayOutputStream bout;
   private PrintStream out;
   private PrintStream originalPrintStream;
@@ -79,13 +79,9 @@ public void setUp() {
   public void tearDown()
       throws InterruptedException, ExecutionException, IOException, TimeoutException {
 
-    try {
-      // Delete the featurestore
-      DeleteFeaturestoreSample.deleteFeaturestoreSample(PROJECT_ID, featurestoreId, USE_FORCE,
-          LOCATION, ENDPOINT, TIMEOUT);
-    } catch (NullPointerException npe) {
-      return;
-    }
+    // Delete the featurestore
+    DeleteFeaturestoreSample.deleteFeaturestoreSample(PROJECT_ID, featurestoreId, USE_FORCE,
+        LOCATION, ENDPOINT, TIMEOUT);
 
     // Assert
     String deleteFeaturestoreResponse = bout.toString();
@@ -94,7 +90,6 @@ public void tearDown()
     System.setOut(originalPrintStream);
   }
 
-  @Ignore
   @Test
   public void testCreateFeaturestoreSample()
       throws IOException, InterruptedException, ExecutionException, TimeoutException {
