feat: Adding tests to default values samples.

m-strzelczyk · Jun 29, 2021 · Jun 15, 2021 · Jun 17, 2021 · Jun 22, 2021 · Jun 25, 2021
commit 7e6674985816387f20316a7c3d7ab96c254d0a1d
diff --git a/samples/snippets/requirements-test.txt b/samples/snippets/requirements-test.txt
@@ -1 +1,2 @@
-pytest==6.2.4
+pytest==6.2.4
+google-cloud-storage==1.39.0
diff --git a/samples/snippets/sample_default_values.py b/samples/snippets/sample_default_values.py
@@ -55,8 +55,11 @@ def set_usage_export_bucket(project_id: str, bucket_name: str,
               "to have the default prefix of `usage_gce`.")
 
     projects_client = compute_v1.ProjectsClient()
-    projects_client.set_usage_export_bucket(
+    operation = projects_client.set_usage_export_bucket(
         project=project_id, usage_export_location_resource=usage_export_location)
+
+    op_client = compute_v1.GlobalOperationsClient()
+    op_client.wait(project=project_id, operation=operation.name)
     return
 # [END compute_usage_report_set]
 
@@ -82,6 +85,7 @@ def get_usage_export_bucket(project_id: str) -> compute_v1.UsageExportLocation:
     if not uel.bucket_name:
         # The Usage Reports are disabled.
         return uel
+
     if not uel.report_name_prefix:
         # Although the server sent the empty value, the next usage report
         # generated with these settings still has the default prefix value
@@ -106,7 +110,10 @@ def disable_usage_export(project_id: str) -> None:
 
     # Updating the setting with None will disable the
     # usage report generation.
-    projects_client.set_usage_export_bucket(
+    operation = projects_client.set_usage_export_bucket(
         project=project_id, usage_export_location_resource=None)
+
+    op_client = compute_v1.GlobalOperationsClient()
+    op_client.wait(project=project_id, operation=operation.name)
     return
 # [END compute_usage_report_disable]
diff --git a/samples/snippets/test_sample_default_values.py b/samples/snippets/test_sample_default_values.py
@@ -0,0 +1,65 @@
+# Copyright 2021 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.
+import typing
+import uuid
+
+import google.auth
+import google.cloud.storage as storage
+import pytest
+
+from sample_default_values import \
+    disable_usage_export, get_usage_export_bucket, set_usage_export_bucket
+
+PROJECT = google.auth.default()[1]
+BUCKET_NAME = "test" + uuid.uuid4().hex[:10]
+TEST_PREFIX = 'some-prefix'
+
+
+@pytest.fixture
+def temp_bucket():
+    storage_client = storage.Client()
+    bucket = storage_client.create_bucket(BUCKET_NAME)
+    yield bucket
+    bucket.delete(force=True)
+
+
+def test_set_usage_export_bucket_default(capsys: typing.Any,
+                                         temp_bucket: storage.Bucket) -> None:
+    set_usage_export_bucket(project_id=PROJECT, bucket_name=temp_bucket.name)
+    uel = get_usage_export_bucket(project_id=PROJECT)
+    assert(uel.bucket_name == temp_bucket.name)
+    assert(uel.report_name_prefix == 'usage_gce')
+    out, _ = capsys.readouterr()
+    assert('default prefix of `usage_gce`.' in out)
+
+    disable_usage_export(project_id=PROJECT)
+    uel = get_usage_export_bucket(project_id=PROJECT)
+    assert(uel.bucket_name == '')
+    assert(uel.report_name_prefix == '')
+
+
+def test_set_usage_export_bucket_custom(capsys: typing.Any,
+                                        temp_bucket: storage.Bucket) -> None:
+    set_usage_export_bucket(project_id=PROJECT, bucket_name=temp_bucket.name,
+                            report_name_prefix=TEST_PREFIX)
+    uel = get_usage_export_bucket(project_id=PROJECT)
+    assert(uel.bucket_name == temp_bucket.name)
+    assert(uel.report_name_prefix == TEST_PREFIX)
+    out, _ = capsys.readouterr()
+    assert('usage_gce' not in out)
+
+    disable_usage_export(project_id=PROJECT)
+    uel = get_usage_export_bucket(project_id=PROJECT)
+    assert(uel.bucket_name == '')
+    assert(uel.report_name_prefix == '')