From 1037f47ef34fa1576402e5391cf62910468e31c9 Mon Sep 17 00:00:00 2001
From: Noah Negrey <nnegrey@users.noreply.github.com>
Date: Thu, 19 Mar 2020 17:32:24 -0600
Subject: [PATCH] automl: move samples into beta set (#3044)

---
 automl/beta/batch_predict.py       | 20 ++++++++++----------
 automl/beta/delete_dataset.py      | 13 +++++--------
 automl/beta/delete_dataset_test.py |  8 ++++----
 3 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/automl/beta/batch_predict.py b/automl/beta/batch_predict.py
index 7d634d2e0a68..911eec733cbb 100644
--- a/automl/beta/batch_predict.py
+++ b/automl/beta/batch_predict.py
@@ -13,17 +13,17 @@
 # limitations under the License.
 
 
-def batch_predict(project_id, model_id, input_uri, output_uri):
-    """Batch predict"""
-    # [START automl_batch_predict_beta]
-    from google.cloud import automl_v1beta1 as automl
+# [START automl_batch_predict_beta]
+from google.cloud import automl_v1beta1 as automl
 
-    # TODO(developer): Uncomment and set the following variables
-    # project_id = "YOUR_PROJECT_ID"
-    # model_id = "YOUR_MODEL_ID"
-    # input_uri = "gs://YOUR_BUCKET_ID/path/to/your/input/csv_or_jsonl"
-    # output_uri = "gs://YOUR_BUCKET_ID/path/to/save/results/"
 
+def batch_predict(
+    project_id="YOUR_PROJECT_ID",
+    model_id="YOUR_MODEL_ID",
+    input_uri="gs://YOUR_BUCKET_ID/path/to/your/input/csv_or_jsonl",
+    output_uri="gs://YOUR_BUCKET_ID/path/to/save/results/",
+):
+    """Batch predict"""
     prediction_client = automl.PredictionServiceClient()
 
     # Get the full path of the model.
@@ -49,4 +49,4 @@ def batch_predict(project_id, model_id, input_uri, output_uri):
             response.result()
         )
     )
-    # [END automl_batch_predict_beta]
+# [END automl_batch_predict_beta]
diff --git a/automl/beta/delete_dataset.py b/automl/beta/delete_dataset.py
index f1935e2b90e4..51647758738d 100644
--- a/automl/beta/delete_dataset.py
+++ b/automl/beta/delete_dataset.py
@@ -13,15 +13,12 @@
 # limitations under the License.
 
 
-def delete_dataset(project_id, dataset_id):
-    """Delete a dataset."""
-    # [START automl_delete_dataset_beta]
-    from google.cloud import automl_v1beta1 as automl
+# [START automl_delete_dataset_beta]
+from google.cloud import automl_v1beta1 as automl
 
-    # TODO(developer): Uncomment and set the following variables
-    # project_id = "YOUR_PROJECT_ID"
-    # dataset_id = "YOUR_DATASET_ID"
 
+def delete_dataset(project_id="YOUR_PROJECT_ID", dataset_id="YOUR_DATASET_ID"):
+    """Delete a dataset."""
     client = automl.AutoMlClient()
     # Get the full path of the dataset
     dataset_full_id = client.dataset_path(
@@ -30,4 +27,4 @@ def delete_dataset(project_id, dataset_id):
     response = client.delete_dataset(dataset_full_id)
 
     print("Dataset deleted. {}".format(response.result()))
-    # [END automl_delete_dataset_beta]
+# [END automl_delete_dataset_beta]
diff --git a/automl/beta/delete_dataset_test.py b/automl/beta/delete_dataset_test.py
index d736aa4df0f4..9781ad260668 100644
--- a/automl/beta/delete_dataset_test.py
+++ b/automl/beta/delete_dataset_test.py
@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import datetime
 import os
+import uuid
 
 from google.cloud import automl_v1beta1 as automl
 import pytest
@@ -28,13 +28,13 @@
 def dataset_id():
     client = automl.AutoMlClient()
     project_location = client.location_path(PROJECT_ID, "us-central1")
-    display_name = "test_" + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+    display_name = "test_{}".format(uuid.uuid4()).replace("-", "")[:32]
     metadata = automl.types.TextExtractionDatasetMetadata()
     dataset = automl.types.Dataset(
         display_name=display_name, text_extraction_dataset_metadata=metadata
     )
-    dataset = client.create_dataset(project_location, dataset)
-    dataset_id = dataset.name.split("/")[-1]
+    response = client.create_dataset(project_location, dataset)
+    dataset_id = response.name.split("/")[-1]
 
     yield dataset_id