samples: automl: create model tests (#1933)

* automl: add create model tests

* Use a fake model to speed up batch predict test

* Update method names and clean up Translate model test

* Rename translate model

* Run code formatter

* Import order

* copy paste typo

* License year

* Fix sample error, fix typos in test

* lint: line length

automl/snippets/src/main/java/com/example/automl/VisionClassificationCreateModel.java

@@ -49,10 +49,7 @@ static void createModel(String projectId, String datasetId, String displayName)
       LocationName projectLocation = LocationName.of(projectId, "us-central1");
       // Set model metadata.
       ImageClassificationModelMetadata metadata =
-          ImageClassificationModelMetadata.newBuilder()
-              .setTrainBudgetMilliNodeHours(
-                  8) // The train budget of creating this model, expressed in hours.
-              .build();
+          ImageClassificationModelMetadata.newBuilder().setTrainBudgetMilliNodeHours(24000).build();
       Model model =
           Model.newBuilder()
               .setDisplayName(displayName)

automl/snippets/src/test/java/com/example/automl/BatchPredictTest.java

@@ -44,7 +44,7 @@
 public class BatchPredictTest {
   private static final String PROJECT_ID = System.getenv("AUTOML_PROJECT_ID");
   private static final String BUCKET_ID = PROJECT_ID + "-lcm";
-  private static final String MODEL_ID = System.getenv("ENTITY_EXTRACTION_MODEL_ID");
+  private static final String MODEL_ID = "TEN0000000000000000000";
   private ByteArrayOutputStream bout;
   private PrintStream out;
 
@@ -62,59 +62,31 @@ public static void checkRequirements() {
   }
 
   @Before
-  public void setUp() throws IOException, ExecutionException, InterruptedException {
-    // Verify that the model is deployed for prediction
-    try (AutoMlClient client = AutoMlClient.create()) {
-      ModelName modelFullId = ModelName.of(PROJECT_ID, "us-central1", MODEL_ID);
-      Model model = client.getModel(modelFullId);
-      if (model.getDeploymentState() == Model.DeploymentState.UNDEPLOYED) {
-        // Deploy the model if not deployed
-        DeployModelRequest request =
-            DeployModelRequest.newBuilder().setName(modelFullId.toString()).build();
-        client.deployModelAsync(request).get();
-      }
-    }
-
+  public void setUp() {
     bout = new ByteArrayOutputStream();
     out = new PrintStream(bout);
     System.setOut(out);
   }
 
   @After
   public void tearDown() {
-    // Delete the created files from GCS
-    Storage storage = StorageOptions.getDefaultInstance().getService();
-    Page<Blob> blobs =
-        storage.list(
-            BUCKET_ID,
-            Storage.BlobListOption.currentDirectory(),
-            Storage.BlobListOption.prefix("TEST_BATCH_PREDICT/"));
-
-    for (Blob blob : blobs.iterateAll()) {
-      Page<Blob> fileBlobs =
-          storage.list(
-              BUCKET_ID,
-              Storage.BlobListOption.currentDirectory(),
-              Storage.BlobListOption.prefix(blob.getName()));
-      for (Blob fileBlob : fileBlobs.iterateAll()) {
-        if (!fileBlob.isDirectory()) {
-          fileBlob.delete();
-        }
-      }
-    }
-
     System.setOut(null);
   }
 
   @Test
-  public void testBatchPredict() throws IOException, ExecutionException, InterruptedException {
-    String inputUri = String.format("gs://%s/entity-extraction/input.jsonl", BUCKET_ID);
-    String outputUri = String.format("gs://%s/TEST_BATCH_PREDICT/", BUCKET_ID);
-    // Act
-    BatchPredict.batchPredict(PROJECT_ID, MODEL_ID, inputUri, outputUri);
-
-    // Assert
-    String got = bout.toString();
-    assertThat(got).contains("Batch Prediction results saved to specified Cloud Storage bucket");
+  public void testBatchPredict() {
+    // As batch prediction can take a long time. Try to batch predict on a model and confirm that
+    // the model was not found, but other elements of the request were valid.
+    try {
+      String inputUri = String.format("gs://%s/entity-extraction/input.jsonl", BUCKET_ID);
+      String outputUri = String.format("gs://%s/TEST_BATCH_PREDICT/", BUCKET_ID);
+      BatchPredict.batchPredict(PROJECT_ID, MODEL_ID, inputUri, outputUri);
+      String got = bout.toString();
+      assertThat(got)
+          .contains("The model is either not found or not supported for prediction yet.");
+    } catch (IOException | ExecutionException | InterruptedException e) {
+      assertThat(e.getMessage())
+          .contains("The model is either not found or not supported for prediction yet.");
+    }
   }
 }

automl/snippets/src/test/java/com/example/automl/LanguageEntityExtractionCreateModelTest.java

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.automl;
+
+import static com.google.common.truth.Truth.assertThat;
+import static junit.framework.TestCase.assertNotNull;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+@SuppressWarnings("checkstyle:abbreviationaswordinname")
+public class LanguageEntityExtractionCreateModelTest {
+
+  private static final String PROJECT_ID = System.getenv("AUTOML_PROJECT_ID");
+  private static final String DATASET_ID = "TEN0000000000000000000";
+  private ByteArrayOutputStream bout;
+  private PrintStream out;
+
+  private static void requireEnvVar(String varName) {
+    assertNotNull(
+        System.getenv(varName),
+        "Environment variable '%s' is required to perform these tests.".format(varName));
+  }
+
+  @BeforeClass
+  public static void checkRequirements() {
+    requireEnvVar("GOOGLE_APPLICATION_CREDENTIALS");
+    requireEnvVar("AUTOML_PROJECT_ID");
+  }
+
+  @Before
+  public void setUp() {
+    bout = new ByteArrayOutputStream();
+    out = new PrintStream(bout);
+    System.setOut(out);
+  }
+
+  @After
+  public void tearDown() {
+    System.setOut(null);
+  }
+
+  @Test
+  public void testLanguageEntityExtractionCreateModel() {
+    // As entity extraction does not let you cancel model creation, instead try to create a model
+    // from a nonexistent dataset, but other elements of the request were valid.
+    try {
+      // Create a random dataset name with a length of 32 characters (max allowed by AutoML)
+      // To prevent name collisions when running tests in multiple java versions at once.
+      // AutoML doesn't allow "-", but accepts "_"
+      String modelName =
+          String.format("test_%s", UUID.randomUUID().toString().replace("-", "_").substring(0, 26));
+      LanguageEntityExtractionCreateModel.createModel(PROJECT_ID, DATASET_ID, modelName);
+      String got = bout.toString();
+      assertThat(got).contains("Dataset does not exist");
+    } catch (IOException | ExecutionException | InterruptedException e) {
+      assertThat(e.getMessage()).contains("Dataset does not exist");
+    }
+  }
+}

automl/snippets/src/test/java/com/example/automl/LanguageSentimentAnalysisCreateModelTest.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.automl;
+
+import static com.google.common.truth.Truth.assertThat;
+import static junit.framework.TestCase.assertNotNull;
+
+import com.google.cloud.automl.v1.AutoMlClient;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+@SuppressWarnings("checkstyle:abbreviationaswordinname")
+public class LanguageSentimentAnalysisCreateModelTest {
+
+  private static final String PROJECT_ID = System.getenv("AUTOML_PROJECT_ID");
+  private static final String DATASET_ID = System.getenv("SENTIMENT_ANALYSIS_DATASET_ID");
+  private ByteArrayOutputStream bout;
+  private PrintStream out;
+  private String operationId;
+
+  private static void requireEnvVar(String varName) {
+    assertNotNull(
+        System.getenv(varName),
+        "Environment variable '%s' is required to perform these tests.".format(varName));
+  }
+
+  @BeforeClass
+  public static void checkRequirements() {
+    requireEnvVar("GOOGLE_APPLICATION_CREDENTIALS");
+    requireEnvVar("AUTOML_PROJECT_ID");
+    requireEnvVar("SENTIMENT_ANALYSIS_DATASET_ID");
+  }
+
+  @Before
+  public void setUp() {
+    bout = new ByteArrayOutputStream();
+    out = new PrintStream(bout);
+    System.setOut(out);
+  }
+
+  @After
+  public void tearDown() throws IOException {
+    // Cancel the operation
+    try (AutoMlClient client = AutoMlClient.create()) {
+      client.getOperationsClient().cancelOperation(operationId);
+    }
+
+    System.setOut(null);
+  }
+
+  @Test
+  public void testLanguageSentimentAnalysisCreateModel()
+      throws IOException, ExecutionException, InterruptedException {
+    // Create a random dataset name with a length of 32 characters (max allowed by AutoML)
+    // To prevent name collisions when running tests in multiple java versions at once.
+    // AutoML doesn't allow "-", but accepts "_"
+    String modelName =
+        String.format("test_%s", UUID.randomUUID().toString().replace("-", "_").substring(0, 26));
+    LanguageSentimentAnalysisCreateModel.createModel(PROJECT_ID, DATASET_ID, modelName);
+
+    String got = bout.toString();
+    assertThat(got).contains("Training started");
+
+    operationId = got.split("Training operation name: ")[1].split("\n")[0];
+  }
+}

automl/snippets/src/test/java/com/example/automl/LanguageTextClassificationCreateModelTest.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.automl;
+
+import static com.google.common.truth.Truth.assertThat;
+import static junit.framework.TestCase.assertNotNull;
+
+import com.google.cloud.automl.v1.AutoMlClient;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+@SuppressWarnings("checkstyle:abbreviationaswordinname")
+public class LanguageTextClassificationCreateModelTest {
+
+  private static final String PROJECT_ID = System.getenv("AUTOML_PROJECT_ID");
+  private static final String DATASET_ID = System.getenv("TEXT_CLASSIFICATION_DATASET_ID");
+  private ByteArrayOutputStream bout;
+  private PrintStream out;
+  private String operationId;
+
+  private static void requireEnvVar(String varName) {
+    assertNotNull(
+        System.getenv(varName),
+        "Environment variable '%s' is required to perform these tests.".format(varName));
+  }
+
+  @BeforeClass
+  public static void checkRequirements() {
+    requireEnvVar("GOOGLE_APPLICATION_CREDENTIALS");
+    requireEnvVar("AUTOML_PROJECT_ID");
+    requireEnvVar("TEXT_CLASSIFICATION_DATASET_ID");
+  }
+
+  @Before
+  public void setUp() {
+    bout = new ByteArrayOutputStream();
+    out = new PrintStream(bout);
+    System.setOut(out);
+  }
+
+  @After
+  public void tearDown() throws IOException {
+    // Cancel the operation
+    try (AutoMlClient client = AutoMlClient.create()) {
+      client.getOperationsClient().cancelOperation(operationId);
+    }
+
+    System.setOut(null);
+  }
+
+  @Test
+  public void testLanguageTextClassificationCreateModel()
+      throws IOException, ExecutionException, InterruptedException {
+    // Create a random dataset name with a length of 32 characters (max allowed by AutoML)
+    // To prevent name collisions when running tests in multiple java versions at once.
+    // AutoML doesn't allow "-", but accepts "_"
+    String modelName =
+        String.format("test_%s", UUID.randomUUID().toString().replace("-", "_").substring(0, 26));
+    LanguageTextClassificationCreateModel.createModel(PROJECT_ID, DATASET_ID, modelName);
+
+    String got = bout.toString();
+    assertThat(got).contains("Training started");
+
+    operationId = got.split("Training operation name: ")[1].split("\n")[0];
+  }
+}