feat(samples): updated the feature samples with close method call and…

… separate timeouts

samples/snippets/src/main/java/aiplatform/CreateFeatureSample.java

@@ -46,16 +46,16 @@ public static void main(String[] args)
     String entityTypeId = "YOUR_ENTITY_TYPE_ID";
     String featureId = "YOUR_FEATURE_ID";
     String description = "YOUR_FEATURE_DESCRIPTION";
-    String valueType = "YOUR_FEATURE_VALUE_TYPE";
+    ValueType valueType = ValueType.STRING;
     String location = "us-central1";
     String endpoint = "us-central1-aiplatform.googleapis.com:443";
-    int timeout = 300;
+    int timeout = 900;
     createFeatureSample(project, featurestoreId, entityTypeId, featureId, description, valueType,
         location, endpoint, timeout);
   }
 
   static void createFeatureSample(String project, String featurestoreId, String entityTypeId,
-      String featureId, String description, String valueType, String location, String endpoint,
+      String featureId, String description, ValueType valueType, String location, String endpoint,
       int timeout) throws IOException, InterruptedException, ExecutionException, TimeoutException {
 
     FeaturestoreServiceSettings featurestoreServiceSettings =
@@ -68,8 +68,7 @@ static void createFeatureSample(String project, String featurestoreId, String en
         FeaturestoreServiceClient.create(featurestoreServiceSettings)) {
 
       Feature feature = Feature.newBuilder().setDescription(description)
-          .setValueType(ValueType.valueOf(valueType))
-          // .setDisableMonitoring(disableMonitoring)
+          .setValueType(valueType)
           .build();
 
       CreateFeatureRequest createFeatureRequest = CreateFeatureRequest.newBuilder()
@@ -83,6 +82,7 @@ static void createFeatureSample(String project, String featurestoreId, String en
       Feature featureResponse = featureFuture.get(timeout, TimeUnit.SECONDS);
       System.out.println("Create Feature Response");
       System.out.format("Name: %s%n", featureResponse.getName());
+      featurestoreServiceClient.close();
     }
   }
 }

samples/snippets/src/main/java/aiplatform/DeleteFeatureSample.java

@@ -21,7 +21,7 @@
 
 package aiplatform;
 
-// [START aiplatform_delete_Feature_sample]
+// [START aiplatform_delete_feature_sample]
 
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.aiplatform.v1.DeleteFeatureRequest;
@@ -74,10 +74,10 @@ static void deleteFeatureSample(String project, String featurestoreId, String en
       System.out.format("Operation name: %s%n", operationFuture.getInitialFuture().get().getName());
       System.out.println("Waiting for operation to finish...");
       operationFuture.get(timeout, TimeUnit.SECONDS);
-
       System.out.format("Deleted Feature.");
+      featurestoreServiceClient.close();
     }
   }
 }
-// [END aiplatform_delete_Feature_sample]
+// [END aiplatform_delete_feature_sample]
 

samples/snippets/src/main/java/aiplatform/GetFeatureSample.java

@@ -64,6 +64,7 @@ static void getFeatureSample(String project, String featurestoreId, String entit
       Feature feature = featurestoreServiceClient.getFeature(getFeatureRequest);
       System.out.println("Get Feature Response");
       System.out.println(feature);
+      featurestoreServiceClient.close();
     }
   }
 }

samples/snippets/src/main/java/aiplatform/ListFeaturesAsyncSample.java

@@ -73,6 +73,7 @@ static void listFeaturesAsyncSample(String project, String featurestoreId, Strin
           break;
         }
       }
+      featurestoreServiceClient.close();
     }
   }
 }

samples/snippets/src/main/java/aiplatform/ListFeaturesSample.java

@@ -62,6 +62,7 @@ static void listFeaturesSample(String project, String featurestoreId, String ent
           .iterateAll()) {
         System.out.println(element);
       }
+      featurestoreServiceClient.close();
     }
   }
 }

samples/snippets/src/main/java/aiplatform/SearchFeaturesAsyncSample.java

@@ -72,6 +72,7 @@ static void searchFeaturesAsyncSample(String project, String query, String locat
           break;
         }
       }
+      featurestoreServiceClient.close();
     }
   }
 }

samples/snippets/src/main/java/aiplatform/SearchFeaturesSample.java

@@ -59,6 +59,7 @@ static void searchFeaturesSample(String project, String query, String location,
           .iterateAll()) {
         System.out.println(element);
       }
+      featurestoreServiceClient.close();
     }
   }
 }

samples/snippets/src/main/java/aiplatform/UpdateFeatureSample.java

@@ -63,6 +63,7 @@ static void updateFeatureSample(String project, String featurestoreId, String en
       Feature featureResponse = featurestoreServiceClient.updateFeature(request);
       System.out.println("Update Feature Response");
       System.out.format("Name: %s%n", featureResponse.getName());
+      featurestoreServiceClient.close();
     }
   }
 }

samples/snippets/src/test/java/aiplatform/FeatureSamplesTest.java

@@ -19,19 +19,12 @@
 import static com.google.common.truth.Truth.assertThat;
 import static junit.framework.TestCase.assertNotNull;
 
-import com.google.api.gax.longrunning.OperationFuture;
-import com.google.cloud.aiplatform.v1.CreateEntityTypeOperationMetadata;
-import com.google.cloud.aiplatform.v1.CreateEntityTypeRequest;
-import com.google.cloud.aiplatform.v1.EntityType;
-import com.google.cloud.aiplatform.v1.FeaturestoreName;
-import com.google.cloud.aiplatform.v1.FeaturestoreServiceClient;
-import com.google.cloud.aiplatform.v1.FeaturestoreServiceSettings;
+import com.google.cloud.aiplatform.v1.Feature.ValueType;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import org.junit.After;
 import org.junit.Before;
@@ -46,13 +39,12 @@ public class FeatureSamplesTest {
   private static final String PROJECT_ID = System.getenv("UCAIP_PROJECT_ID");
   private static final int MIN_NODE_COUNT = 1;
   private static final int MAX_NODE_COUNT = 5;
-  private static final String VALUE_TYPE = "STRING";
+  private static final ValueType VALUE_TYPE = ValueType.STRING;
   private static final String DESCRIPTION = "Test Description";
   private static final boolean USE_FORCE = true;
   private static final String QUERY = "value_type=STRING";
   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 ByteArrayOutputStream bout;
   private PrintStream out;
   private PrintStream originalPrintStream;
@@ -77,44 +69,14 @@ public void setUp() {
     originalPrintStream = System.out;
     System.setOut(out);
   }
-
-  static void createEntityTypeSample(String project, String featurestoreId, String entityTypeId,
-      String description, String location, String endpoint, int timeout)
-      throws IOException, InterruptedException, ExecutionException, TimeoutException {
-
-    FeaturestoreServiceSettings featurestoreServiceSettings =
-        FeaturestoreServiceSettings.newBuilder().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
-    // the "close" method on the client to safely clean up any remaining background resources.
-    try (FeaturestoreServiceClient featurestoreServiceClient =
-        FeaturestoreServiceClient.create(featurestoreServiceSettings)) {
-
-      EntityType entityType = EntityType.newBuilder().setDescription(description).build();
-
-      CreateEntityTypeRequest createEntityTypeRequest = CreateEntityTypeRequest.newBuilder()
-          .setParent(FeaturestoreName.of(project, location, featurestoreId).toString())
-          .setEntityType(entityType).setEntityTypeId(entityTypeId).build();
-
-      OperationFuture<EntityType, CreateEntityTypeOperationMetadata> entityTypeFuture =
-          featurestoreServiceClient.createEntityTypeAsync(createEntityTypeRequest);
-      System.out.format("Operation name: %s%n",
-          entityTypeFuture.getInitialFuture().get().getName());
-      System.out.println("Waiting for operation to finish...");
-      EntityType entityTypeResponse = entityTypeFuture.get(timeout, TimeUnit.SECONDS);
-      System.out.println("Create Entity Type Response");
-      System.out.format("Name: %s%n", entityTypeResponse.getName());
-    }
-  }
 
   @After
   public void tearDown()
       throws InterruptedException, ExecutionException, IOException, TimeoutException {
 
     // Delete the featurestore
     DeleteFeaturestoreSample.deleteFeaturestoreSample(PROJECT_ID, featurestoreId, USE_FORCE,
-        LOCATION, ENDPOINT, 60);
+        LOCATION, ENDPOINT, 300);
 
     // Assert
     String deleteFeaturestoreResponse = bout.toString();
@@ -143,8 +105,8 @@ public void testFeatureSamples()
     // Create the entity type
     String entityTypeTempUuid = UUID.randomUUID().toString().replaceAll("-", "_").substring(0, 16);
     String entityTypeId = String.format("temp_create_entity_type_test_%s", entityTypeTempUuid);
-    createEntityTypeSample(PROJECT_ID, featurestoreId, entityTypeId,
-        DESCRIPTION, LOCATION, ENDPOINT, TIMEOUT);
+    CreateEntityTypeSample.createEntityTypeSample(PROJECT_ID, featurestoreId, entityTypeId,
+        DESCRIPTION, LOCATION, ENDPOINT, 900);
 
     // Assert
     String createEntityTypeResponse = bout.toString();
@@ -154,7 +116,7 @@ public void testFeatureSamples()
     String featureTempUuid = UUID.randomUUID().toString().replaceAll("-", "_").substring(0, 26);
     String featureId = String.format("temp_create_feature_test_%s", featureTempUuid);
     CreateFeatureSample.createFeatureSample(PROJECT_ID, featurestoreId, entityTypeId, featureId,
-        DESCRIPTION, VALUE_TYPE, LOCATION, ENDPOINT, TIMEOUT);
+        DESCRIPTION, VALUE_TYPE, LOCATION, ENDPOINT, 900);
 
     // Assert
     String createFeatureResponse = bout.toString();
@@ -200,7 +162,7 @@ public void testFeatureSamples()
 
     // Delete the feature
     DeleteFeatureSample.deleteFeatureSample(PROJECT_ID, featurestoreId, entityTypeId, featureId,
-        LOCATION, ENDPOINT, TIMEOUT);
+        LOCATION, ENDPOINT, 300);
 
     // Assert
     String deleteFeatureResponse = bout.toString();