Add consistency group clone sample + updated tests

Author: Thoughtseize1

compute/client_library/ingredients/disks/сonsistency_groups/clone_disks_consistency_group.py

@@ -0,0 +1,52 @@
+#  Copyright 2024 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.
+
+# This is an ingredient file. It is not meant to be run directly. Check the samples/snippets
+# folder for complete code samples that are ready to be used.
+# Disabling flake8 for the ingredients file, as it would fail F821 - undefined name check.
+# flake8: noqa
+
+from google.cloud import compute_v1
+
+
+# <INGREDIENT consistency_group_clone_disks>
+def clone_disks_to_consistency_group(project_id, group_name, group_region):
+    """
+    Clones disks to a consistency group in the specified region.
+    Args:
+        project_id (str): The ID of the Google Cloud project.
+        group_name (str): The name of the consistency group.
+        group_region (str): The region where the consistency group is located.
+    Returns:
+        bool: True if the disks were successfully cloned to the consistency group.
+    """
+    consistency_group_policy = (
+        f"projects/{project_id}/regions/{group_region}/resourcePolicies/{group_name}"
+    )
+
+    resource = compute_v1.BulkInsertDiskResource(
+        source_consistency_group_policy=consistency_group_policy
+    )
+    client = compute_v1.RegionDisksClient()
+    request = compute_v1.BulkInsertRegionDiskRequest(
+        project=project_id,
+        region=group_region,
+        bulk_insert_disk_resource_resource=resource,
+    )
+    operation = client.bulk_insert(request=request)
+    wait_for_extended_operation(operation, verbose_name="bulk insert disk")
+    return True
+
+
+# </INGREDIENT>

compute/client_library/recipes/disks/сonsistency_groups/clone_disks_consistency_group.py

@@ -0,0 +1,23 @@
+#  Copyright 2024 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.
+# flake8: noqa
+
+# <REGION compute_consistency_group_clone>
+# <IMPORTS/>
+
+# <INGREDIENT wait_for_extended_operation />
+
+# <INGREDIENT consistency_group_clone_disks />
+
+# </REGION compute_consistency_group_clone>

compute/client_library/snippets/disks/сonsistency_groups/clone_disks_consistency_group.py

@@ -0,0 +1,108 @@
+#  Copyright 2024 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.
+# flake8: noqa
+
+
+# This file is automatically generated. Please do not modify it directly.
+# Find the relevant recipe file in the samples/recipes or samples/ingredients
+# directory and apply your changes there.
+
+
+# [START compute_consistency_group_clone]
+from __future__ import annotations
+
+import sys
+from typing import Any
+
+from google.api_core.extended_operation import ExtendedOperation
+from google.cloud import compute_v1
+
+
+def wait_for_extended_operation(
+    operation: ExtendedOperation, verbose_name: str = "operation", timeout: int = 300
+) -> Any:
+    """
+    Waits for the extended (long-running) operation to complete.
+
+    If the operation is successful, it will return its result.
+    If the operation ends with an error, an exception will be raised.
+    If there were any warnings during the execution of the operation
+    they will be printed to sys.stderr.
+
+    Args:
+        operation: a long-running operation you want to wait on.
+        verbose_name: (optional) a more verbose name of the operation,
+            used only during error and warning reporting.
+        timeout: how long (in seconds) to wait for operation to finish.
+            If None, wait indefinitely.
+
+    Returns:
+        Whatever the operation.result() returns.
+
+    Raises:
+        This method will raise the exception received from `operation.exception()`
+        or RuntimeError if there is no exception set, but there is an `error_code`
+        set for the `operation`.
+
+        In case of an operation taking longer than `timeout` seconds to complete,
+        a `concurrent.futures.TimeoutError` will be raised.
+    """
+    result = operation.result(timeout=timeout)
+
+    if operation.error_code:
+        print(
+            f"Error during {verbose_name}: [Code: {operation.error_code}]: {operation.error_message}",
+            file=sys.stderr,
+            flush=True,
+        )
+        print(f"Operation ID: {operation.name}", file=sys.stderr, flush=True)
+        raise operation.exception() or RuntimeError(operation.error_message)
+
+    if operation.warnings:
+        print(f"Warnings during {verbose_name}:\n", file=sys.stderr, flush=True)
+        for warning in operation.warnings:
+            print(f" - {warning.code}: {warning.message}", file=sys.stderr, flush=True)
+
+    return result
+
+
+def clone_disks_to_consistency_group(project_id, group_name, group_region):
+    """
+    Clones disks to a consistency group in the specified region.
+    Args:
+        project_id (str): The ID of the Google Cloud project.
+        group_name (str): The name of the consistency group.
+        group_region (str): The region where the consistency group is located.
+    Returns:
+        bool: True if the disks were successfully cloned to the consistency group.
+    """
+    consistency_group_policy = (
+        f"projects/{project_id}/regions/{group_region}/resourcePolicies/{group_name}"
+    )
+
+    resource = compute_v1.BulkInsertDiskResource(
+        source_consistency_group_policy=consistency_group_policy
+    )
+    client = compute_v1.RegionDisksClient()
+    request = compute_v1.BulkInsertRegionDiskRequest(
+        project=project_id,
+        region=group_region,
+        bulk_insert_disk_resource_resource=resource,
+    )
+    operation = client.bulk_insert(request=request)
+    wait_for_extended_operation(operation, verbose_name="bulk insert disk")
+    return True
+
+
+# [END compute_consistency_group_clone]

compute/client_library/snippets/tests/test_consistency_groups.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 import os
 
+from unittest.mock import patch
+
 import uuid
 
 import pytest
@@ -21,6 +23,9 @@
 from ..disks.сonsistency_groups.add_disk_consistency_group import (
     add_disk_consistency_group,
 )
+from ..disks.сonsistency_groups.clone_disks_consistency_group import (
+    clone_disks_to_consistency_group,
+)
 from ..disks.сonsistency_groups.create_consistency_group import create_consistency_group
 from ..disks.сonsistency_groups.delete_consistency_group import delete_consistency_group
 from ..disks.сonsistency_groups.list_disks_consistency_group import (
@@ -94,3 +99,10 @@ def test_add_remove_and_list_disks_consistency_group(
         consistency_group_region=REGION,
     )
     assert not disks
+
+
+def test_clone_disk_consistency_group():
+    with patch("google.cloud.compute_v1.RegionDisksClient") as mock_client:
+        client = mock_client.return_value
+        clone_disks_to_consistency_group(PROJECT_ID, "group-4", "us-east1")
+        client.bulk_insert.assert_called_once()

compute/client_library/snippets/tests/test_disks.py

@@ -38,6 +38,15 @@
 from ..disks.replication_disk_start import start_disk_replication
 from ..disks.replication_disk_stop import stop_disk_replication
 from ..disks.resize_disk import resize_disk
+from ..disks.сonsistency_groups.add_disk_consistency_group import (
+    add_disk_consistency_group,
+)
+from ..disks.сonsistency_groups.clone_disks_consistency_group import (
+    clone_disks_to_consistency_group,
+)
+from ..disks.сonsistency_groups.create_consistency_group import create_consistency_group
+from ..disks.сonsistency_groups.delete_consistency_group import delete_consistency_group
+from ..disks.сonsistency_groups.remove_disk_consistency_group import remove_disk_consistency_group
 from ..images.get import get_image_from_family
 from ..instances.create import create_instance, disk_from_image
 from ..instances.delete import delete_instance
@@ -50,7 +59,7 @@
 ZONE = "europe-west2-c"
 ZONE_SECONDARY = "europe-west1-c"
 REGION = "europe-west2"
-REGION_SECONDARY = "europe-central2"
+REGION_SECONDARY = "europe-west3"
 KMS_KEYRING_NAME = "compute-test-keyring"
 KMS_KEY_NAME = "compute-test-key"
 DISK_SIZE = 11
@@ -499,3 +508,69 @@ def test_start_stop_zone_replication(test_empty_pd_balanced_disk, autodelete_dis
     )
     # Wait for the replication to stop
     time.sleep(20)
+
+
+def test_clone_disks_in_consistency_group(
+    autodelete_regional_disk_name,
+    autodelete_regional_blank_disk,
+):
+    first = "first-group" + uuid.uuid4().hex[:5]
+    create_consistency_group(PROJECT, REGION, first, "description")
+
+    add_disk_consistency_group(
+        project_id=PROJECT,
+        disk_name=autodelete_regional_blank_disk.name,
+        disk_location=REGION,
+        consistency_group_name=first,
+        consistency_group_region=REGION,
+    )
+
+    second_disk = create_secondary_region_disk(
+        autodelete_regional_blank_disk.name,
+        PROJECT,
+        REGION,
+        autodelete_regional_disk_name,
+        PROJECT,
+        REGION_SECONDARY,
+        DISK_SIZE,
+    )
+
+    second = "second-group" + uuid.uuid4().hex[:5]
+    create_consistency_group(PROJECT, REGION_SECONDARY, second, "description")
+
+    add_disk_consistency_group(
+        project_id=PROJECT,
+        disk_name=second_disk.name,
+        disk_location=REGION_SECONDARY,
+        consistency_group_name=second,
+        consistency_group_region=REGION_SECONDARY,
+    )
+
+    start_disk_replication(
+        project_id=PROJECT,
+        primary_disk_location=REGION,
+        primary_disk_name=autodelete_regional_blank_disk.name,
+        secondary_disk_location=REGION_SECONDARY,
+        secondary_disk_name=autodelete_regional_disk_name,
+    )
+    time.sleep(45)
+    try:
+        assert clone_disks_to_consistency_group(PROJECT, second, REGION_SECONDARY)
+    finally:
+        stop_disk_replication(
+            project_id=PROJECT,
+            primary_disk_location=REGION,
+            primary_disk_name=autodelete_regional_blank_disk.name,
+        )
+        # Wait for the replication to stop
+        time.sleep(30)
+        disks = compute_v1.RegionDisksClient().list(
+            project=PROJECT, region=REGION_SECONDARY
+        )
+        if disks:
+            for disk in disks:
+                delete_regional_disk(PROJECT, REGION_SECONDARY, disk.name)
+        time.sleep(20)
+        remove_disk_consistency_group(PROJECT, autodelete_regional_blank_disk.name, REGION, first, REGION)
+        delete_consistency_group(PROJECT, REGION, first)
+        delete_consistency_group(PROJECT, REGION_SECONDARY, second)