automl: move samples into beta set

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]

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]

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