feat: adding samples to start/stop/reset operations (#75)

m-strzelczyk · parthea · web-flow · commit 4b841efe4cb8 · 2021-08-04T12:14:01.000+02:00
* feat: Adding start/stop compute samples.

Co-authored-by: Anthonios Partheniou &lt;partheniou@google.com&gt;
diff --git a/packages/google-cloud-compute/samples/snippets/sample_start_stop.py b/packages/google-cloud-compute/samples/snippets/sample_start_stop.py
@@ -0,0 +1,116 @@
+#!/usr/bin/env python
+
+# 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.
+
+"""
+A sample script showing how to start and stop Google Compute Engine instances.
+"""
+from google.cloud import compute_v1
+
+
+# [START compute_start_instance]
+def start_instance(project_id: str, zone: str, instance_name: str):
+    """
+    Starts a stopped Google Compute Engine instance (with unencrypted disks).
+
+    Args:
+        project_id: project ID or project number of the Cloud project your instance belongs to.
+        zone: name of the zone your instance belongs to.
+        instance_name: name of the instance your want to start.
+    """
+    instance_client = compute_v1.InstancesClient()
+    op_client = compute_v1.ZoneOperationsClient()
+
+    op = instance_client.start(project=project_id, zone=zone, instance=instance_name)
+
+    op_client.wait(project=project_id, zone=zone, operation=op.name)
+    return
+# [END compute_start_instance]
+
+
+# [START compute_start_enc_instance]
+def start_instance_with_encryption_key(project_id: str, zone: str, instance_name: str, key: bytes):
+    """
+    Starts a stopped Google Compute Engine instance (with encrypted disks).
+
+    Args:
+        project_id: project ID or project number of the Cloud project your instance belongs to.
+        zone: name of the zone your instance belongs to.
+        instance_name: name of the instance your want to start.
+        key: bytes object representing a raw base64 encoded key to your machines boot disk.
+            For more information about disk encryption see:
+            https://cloud.google.com/compute/docs/disks/customer-supplied-encryption#specifications
+    """
+    instance_client = compute_v1.InstancesClient()
+    op_client = compute_v1.ZoneOperationsClient()
+
+    instance_data = instance_client.get(project=project_id, zone=zone, instance=instance_name)
+
+    # Prepare the information about disk encryption
+    disk_data = compute_v1.CustomerEncryptionKeyProtectedDisk()
+    disk_data.source = instance_data.disks[0].source
+    disk_data.disk_encryption_key = compute_v1.CustomerEncryptionKey()
+    # Use raw_key to send over the key to unlock the disk
+    # To use a key stored in KMS, you need to provide `kms_key_name` and `kms_key_service_account`
+    disk_data.disk_encryption_key.raw_key = key
+    enc_data = compute_v1.InstancesStartWithEncryptionKeyRequest()
+    enc_data.disks = [disk_data]
+
+    op = instance_client.start_with_encryption_key(project=project_id, zone=zone, instance=instance_name,
+                                                   instances_start_with_encryption_key_request_resource=enc_data)
+
+    op_client.wait(project=project_id, zone=zone, operation=op.name)
+    return
+# [END compute_start_enc_instance]
+
+
+# [START compute_stop_instance]
+def stop_instance(project_id: str, zone: str, instance_name: str):
+    """
+    Stops a stopped Google Compute Engine instance.
+
+    Args:
+        project_id: project ID or project number of the Cloud project your instance belongs to.
+        zone: name of the zone your instance belongs to.
+        instance_name: name of the instance your want to stop.
+    """
+    instance_client = compute_v1.InstancesClient()
+    op_client = compute_v1.ZoneOperationsClient()
+
+    op = instance_client.stop(project=project_id, zone=zone, instance=instance_name)
+
+    op_client.wait(project=project_id, zone=zone, operation=op.name)
+    return
+# [END compute_stop_instance]
+
+
+# [START compute_reset_instance]
+def reset_instance(project_id: str, zone: str, instance_name: str):
+    """
+    Resets a stopped Google Compute Engine instance (with unencrypted disks).
+
+    Args:
+        project_id: project ID or project number of the Cloud project your instance belongs to.
+        zone: name of the zone your instance belongs to.
+        instance_name: name of the instance your want to reset.
+    """
+    instance_client = compute_v1.InstancesClient()
+    op_client = compute_v1.ZoneOperationsClient()
+
+    op = instance_client.reset(project=project_id, zone=zone, instance=instance_name)
+
+    op_client.wait(project=project_id, zone=zone, operation=op.name)
+    return
+# [END compute_reset_instance]
diff --git a/packages/google-cloud-compute/samples/snippets/test_sample_default_values.py b/packages/google-cloud-compute/samples/snippets/test_sample_default_values.py
@@ -11,6 +11,7 @@
 # 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 time
 import typing
 import uuid
 
@@ -37,13 +38,15 @@ def temp_bucket():
 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)
+    time.sleep(5)  # To make sure the settings are properly updated
     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)
+    time.sleep(5)  # To make sure the settings are properly updated
     uel = get_usage_export_bucket(project_id=PROJECT)
     assert(uel.bucket_name == '')
     assert(uel.report_name_prefix == '')
@@ -53,13 +56,15 @@ 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)
+    time.sleep(5)  # To make sure the settings are properly updated
     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)
+    time.sleep(5)  # To make sure the settings are properly updated
     uel = get_usage_export_bucket(project_id=PROJECT)
     assert(uel.bucket_name == '')
     assert(uel.report_name_prefix == '')
diff --git a/packages/google-cloud-compute/samples/snippets/test_sample_start_stop.py b/packages/google-cloud-compute/samples/snippets/test_sample_start_stop.py
@@ -0,0 +1,150 @@
+# 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 base64
+import random
+import string
+import time
+import uuid
+
+import google.auth
+from google.cloud import compute_v1
+
+import pytest
+
+from samples.snippets.sample_start_stop import start_instance, start_instance_with_encryption_key, stop_instance
+
+PROJECT = google.auth.default()[1]
+
+INSTANCE_ZONE = "europe-central2-b"
+
+KEY = "".join(random.sample(string.ascii_letters, 32))
+KEY_B64 = base64.b64encode(KEY.encode())  # for example: b'VEdORldtY3NKellPdWRDcUF5YlNVREtJdm5qaFJYSFA='
+
+
+def _make_disk(raw_key: bytes = None):
+    disk = compute_v1.AttachedDisk()
+    initialize_params = compute_v1.AttachedDiskInitializeParams()
+    initialize_params.source_image = (
+        "projects/debian-cloud/global/images/family/debian-10"
+    )
+    initialize_params.disk_size_gb = "10"
+    disk.initialize_params = initialize_params
+    disk.auto_delete = True
+    disk.boot = True
+    disk.type_ = compute_v1.AttachedDisk.Type.PERSISTENT
+
+    if raw_key:
+        disk.disk_encryption_key = compute_v1.CustomerEncryptionKey()
+        disk.disk_encryption_key.raw_key = raw_key
+
+    return disk
+
+
+def _make_request(disk: compute_v1.AttachedDisk):
+    network_interface = compute_v1.NetworkInterface()
+    network_interface.name = 'default'
+    network_interface.access_configs = []
+
+    # Collect information into the Instance object.
+    instance = compute_v1.Instance()
+    instance.name = "i" + uuid.uuid4().hex[:10]
+    instance.disks = [disk]
+    full_machine_type_name = f"zones/{INSTANCE_ZONE}/machineTypes/e2-micro"
+    instance.machine_type = full_machine_type_name
+    instance.network_interfaces = [network_interface]
+
+    # Prepare the request to insert an instance.
+    request = compute_v1.InsertInstanceRequest()
+    request.zone = INSTANCE_ZONE
+    request.project = PROJECT
+    request.instance_resource = instance
+    return request
+
+
+def _create_instance(request: compute_v1.InsertInstanceRequest):
+    instance_client = compute_v1.InstancesClient()
+    operation_client = compute_v1.ZoneOperationsClient()
+
+    operation = instance_client.insert(request=request)
+    operation_client.wait(operation=operation.name, zone=INSTANCE_ZONE, project=PROJECT)
+
+    return instance_client.get(project=PROJECT, zone=INSTANCE_ZONE, instance=request.instance_resource.name)
+
+
+def _delete_instance(instance: compute_v1.Instance):
+    instance_client = compute_v1.InstancesClient()
+    operation_client = compute_v1.ZoneOperationsClient()
+
+    operation = instance_client.delete(project=PROJECT, zone=INSTANCE_ZONE, instance=instance.name)
+    operation_client.wait(operation=operation.name, zone=INSTANCE_ZONE, project=PROJECT)
+
+
+def _get_status(instance: compute_v1.Instance) -> compute_v1.Instance.Status:
+    instance_client = compute_v1.InstancesClient()
+    return instance_client.get(project=PROJECT, zone=INSTANCE_ZONE, instance=instance.name).status
+
+
+@pytest.fixture
+def compute_instance():
+    disk = _make_disk()
+    request = _make_request(disk)
+
+    instance = _create_instance(request)
+
+    yield instance
+
+    _delete_instance(instance)
+
+
+@pytest.fixture
+def compute_encrypted_instance():
+    disk = _make_disk(KEY_B64)
+    request = _make_request(disk)
+
+    instance = _create_instance(request)
+
+    yield instance
+
+    _delete_instance(instance)
+
+
+def test_instance_operations(compute_instance):
+    assert _get_status(compute_instance) == compute_v1.Instance.Status.RUNNING
+
+    stop_instance(PROJECT, INSTANCE_ZONE, compute_instance.name)
+
+    while _get_status(compute_instance) == compute_v1.Instance.Status.STOPPING:
+        # Since we can't configure timeout parameter for operation wait() (b/188037306)
+        # We need to do some manual waiting for the stopping to finish...
+        time.sleep(5)
+
+    assert _get_status(compute_instance) == compute_v1.Instance.Status.TERMINATED
+
+    start_instance(PROJECT, INSTANCE_ZONE, compute_instance.name)
+    assert _get_status(compute_instance) == compute_v1.Instance.Status.RUNNING
+
+
+def test_instance_encrypted(compute_encrypted_instance):
+    assert _get_status(compute_encrypted_instance) == compute_v1.Instance.Status.RUNNING
+
+    stop_instance(PROJECT, INSTANCE_ZONE, compute_encrypted_instance.name)
+    while _get_status(compute_encrypted_instance) == compute_v1.Instance.Status.STOPPING:
+        # Since we can't configure timeout parameter for operation wait() (b/188037306)
+        # We need to do some manual waiting for the stopping to finish...
+        time.sleep(5)
+
+    assert _get_status(compute_encrypted_instance) == compute_v1.Instance.Status.TERMINATED
+
+    start_instance_with_encryption_key(PROJECT, INSTANCE_ZONE, compute_encrypted_instance.name, KEY_B64)
+    assert _get_status(compute_encrypted_instance) == compute_v1.Instance.Status.RUNNING
