Repository: google/cloud-forensics-utils

Issue:
disk_type in CreateDiskCopy should default to the source's disk type
In CreateDisk(Volume)Copy methods, it'd be nice for the disk copy to have the same disk type as the source disk, instead of a default disk type for all copies.

Discussion:
codecov-commenter: # [Codecov](https://codecov.io/gh/google/cloud-forensics-utils/pull/219?src=pr&el=h1) Report
> Merging [#219](https://codecov.io/gh/google/cloud-forensics-utils/pull/219?src=pr&el=desc) into [master](https://codecov.io/gh/google/cloud-forensics-utils/commit/e4448c60002eee03c344ec2803276ff5d91936be&el=desc) will **increase** coverage by `2.69%`.
> The diff coverage is `71.24%`.

[![Impacted file tree graph](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/graphs/tree.svg?width=650&height=150&src=pr&token=zZkPno3ah0)](https://codecov.io/gh/google/cloud-forensics-utils/pull/219?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #219      +/-   ##
==========================================
+ Coverage   65.39%   68.08%   +2.69%     
==========================================
  Files          14       22       +8     
  Lines         968     1711     +743     
==========================================
+ Hits          633     1165     +532     
- Misses        335      546     +211     
```

| Flag | Coverage Δ | |
|---|---|---|
| #nosetests | `68.08% <71.24%> (+2.69%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/google/cloud-forensics-utils/pull/219?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [libcloudforensics/providers/aws/forensics.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2F3cy9mb3JlbnNpY3MucHk=) | `55.71% <45.45%> (-14.06%)` | :arrow_down: |
| [libcloudforensics/providers/gcp/forensics.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9mb3JlbnNpY3MucHk=) | `54.92% <51.11%> (-10.08%)` | :arrow_down: |
| [...ibcloudforensics/providers/gcp/internal/compute.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9jb21wdXRlLnB5) | `56.56% <51.93%> (-5.80%)` | :arrow_down: |
| [...cs/providers/gcp/internal/compute\_base\_resource.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9jb21wdXRlX2Jhc2VfcmVzb3VyY2UucHk=) | `37.33% <52.94%> (+0.82%)` | :arrow_up: |
| [...cloudforensics/providers/azure/internal/compute.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2F6dXJlL2ludGVybmFsL2NvbXB1dGUucHk=) | `54.34% <54.34%> (ø)` | |
| [libcloudforensics/providers/aws/internal/log.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2F3cy9pbnRlcm5hbC9sb2cucHk=) | `73.07% <55.55%> (-4.20%)` | :arrow_down: |
| [...ibcloudforensics/providers/gcp/internal/project.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9wcm9qZWN0LnB5) | `61.90% <66.66%> (-3.81%)` | :arrow_down: |
| [...bcloudforensics/providers/azure/internal/common.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2F6dXJlL2ludGVybmFsL2NvbW1vbi5weQ==) | `72.05% <72.05%> (ø)` | |
| [libcloudforensics/providers/gcp/internal/gke.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9na2UucHk=) | `73.68% <73.68%> (ø)` | |
| [libcloudforensics/providers/gcp/internal/common.py](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree#diff-bGliY2xvdWRmb3JlbnNpY3MvcHJvdmlkZXJzL2djcC9pbnRlcm5hbC9jb21tb24ucHk=) | `60.00% <75.86%> (+3.67%)` | :arrow_up: |
| ... and [21 more](https://codecov.io/gh/google/cloud-forensics-utils/pull/219/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/google/cloud-forensics-utils/pull/219?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/google/cloud-forensics-utils/pull/219?src=pr&el=footer). Last update [cc650e1...308de10](https://codecov.io/gh/google/cloud-forensics-utils/pull/219?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/libcloudforensics/providers/azure/forensics.py b/libcloudforensics/providers/azure/forensics.py
index de7d1f8..3c03381 100644
--- a/libcloudforensics/providers/azure/forensics.py
+++ b/libcloudforensics/providers/azure/forensics.py
@@ -30,7 +30,7 @@ def CreateDiskCopy(
     resource_group_name: str,
     instance_name: Optional[str] = None,
     disk_name: Optional[str] = None,
-    disk_type: str = 'Standard_LRS',
+    disk_type: Optional[str] = None,
     region: str = 'eastus',
     src_profile: Optional[str] = None,
     dst_profile: Optional[str] = None) -> 'compute.AZDisk':
@@ -47,7 +47,7 @@ def CreateDiskCopy(
         then instance_name needs to be set and the boot disk will be copied.
     disk_type (str): Optional. The sku name for the disk to create. Can be
         Standard_LRS, Premium_LRS, StandardSSD_LRS, or UltraSSD_LRS. The
-        default value is Standard_LRS.
+        default behavior is to use the same disk type as the source disk.
     region (str): Optional. The region in which to create the disk copy.
         Default is eastus.
     src_profile (str): Optional. The name of the source profile to use for the
@@ -88,6 +88,10 @@ def CreateDiskCopy(
       disk_to_copy = instance.GetBootDisk()
     logger.info('Disk copy of {0:s} started...'.format(
         disk_to_copy.name))
+
+    if not disk_type:
+      disk_type = disk_to_copy.GetDiskType()
+
     snapshot = disk_to_copy.Snapshot()
 
     subscription_ids = src_account.ListSubscriptionIDs()
diff --git a/libcloudforensics/providers/azure/internal/account.py b/libcloudforensics/providers/azure/internal/account.py
index 79b37be..3753f0d 100644
--- a/libcloudforensics/providers/azure/internal/account.py
+++ b/libcloudforensics/providers/azure/internal/account.py
@@ -245,7 +245,8 @@ class AZAccount:
         'creation_data': {
             'sourceResourceId': snapshot.resource_id,
             'create_option': models.DiskCreateOption.copy
-        }
+        },
+        'sku': {'name': disk_type}
     }
 
     try:
@@ -253,8 +254,7 @@ class AZAccount:
       request = self.compute_client.disks.create_or_update(
           self.default_resource_group_name,
           disk_name,
-          creation_data,
-          sku=disk_type)
+          creation_data)
       while not request.done():
         sleep(5)  # Wait 5 seconds before checking disk status again
       disk = request.result()
@@ -360,7 +360,8 @@ class AZAccount:
             'source_uri': copied_blob.url,
             'storage_account_id': storage_account_id,
             'create_option': models.DiskCreateOption.import_enum
-        }
+        },
+        'sku': {'name': disk_type}
     }
 
     try:
@@ -368,8 +369,7 @@ class AZAccount:
       request = self.compute_client.disks.create_or_update(
           self.default_resource_group_name,
           disk_name,
-          creation_data,
-          sku=disk_type)
+          creation_data)
       while not request.done():
         sleep(5)  # Wait 5 seconds before checking disk status again
       disk = request.result()
diff --git a/libcloudforensics/providers/azure/internal/compute.py b/libcloudforensics/providers/azure/internal/compute.py
index a343dbf..b9879d3 100644
--- a/libcloudforensics/providers/azure/internal/compute.py
+++ b/libcloudforensics/providers/azure/internal/compute.py
@@ -273,6 +273,17 @@ class AZDisk(AZComputeResource):
                       snapshot.location,
                       self)
 
+  def GetDiskType(self) -> str:
+    """Return the SKU disk type.
+
+    Returns:
+      str: The SKU disk type.
+    """
+    disk = self.az_account.compute_client.disks.get(
+        self.resource_group_name, self.name)
+    disk_type = disk.sku.name  # type: str
+    return disk_type
+
 
 class AZSnapshot(AZComputeResource):
   """Class that represents Azure snapshots.
diff --git a/libcloudforensics/providers/gcp/forensics.py b/libcloudforensics/providers/gcp/forensics.py
index 5b6fdbf..f308bb2 100644
--- a/libcloudforensics/providers/gcp/forensics.py
+++ b/libcloudforensics/providers/gcp/forensics.py
@@ -37,7 +37,7 @@ def CreateDiskCopy(
     zone: str,
     instance_name: Optional[str] = None,
     disk_name: Optional[str] = None,
-    disk_type: str = 'pd-standard') -> 'compute.GoogleComputeDisk':
+    disk_type: Optional[str] = None) -> 'compute.GoogleComputeDisk':
   """Creates a copy of a Google Compute Disk.
 
   Args:
@@ -48,8 +48,8 @@ def CreateDiskCopy(
     disk_name (str): Optional. Name of the disk to copy. If None,
         instance_name must be specified and the boot disk will be copied.
     disk_type (str): Optional. URL of the disk type resource describing
-        which disk type to use to create the disk. Default is pd-standard. Use
-        pd-ssd to have a SSD disk.
+        which disk type to use to create the disk. The default behavior is to
+        use the same disk type as the source disk.
 
   Returns:
     GoogleComputeDisk: A Google Compute Disk object.
@@ -73,6 +73,9 @@ def CreateDiskCopy(
       instance = src_project.compute.GetInstance(instance_name)
       disk_to_copy = instance.GetBootDisk()
 
+    if not disk_type:
+      disk_type = disk_to_copy.GetDiskType()
+
     logger.info('Disk copy of {0:s} started...'.format(
         disk_to_copy.name))
     snapshot = disk_to_copy.Snapshot()
diff --git a/libcloudforensics/providers/gcp/internal/compute.py b/libcloudforensics/providers/gcp/internal/compute.py
index 225a1c6..6c1bc66 100644
--- a/libcloudforensics/providers/gcp/internal/compute.py
+++ b/libcloudforensics/providers/gcp/internal/compute.py
@@ -917,6 +917,17 @@ class GoogleComputeDisk(compute_base_resource.GoogleComputeBaseResource):
     self.BlockOperation(response, zone=self.zone)
     return GoogleComputeSnapshot(disk=self, name=snapshot_name)
 
+  def GetDiskType(self) -> str:
+    """Return the disk type.
+
+    Returns:
+      str: The disk type.
+    """
+    # 'type': https://www.googleapis.com/compute/v1/projects/<>/zones/us
+    # -central1-a/diskTypes/pd-standard
+    disk_type = self.GetOperation()['type'].split('/')[-1]  # type: str
+    return disk_type
+
 
 class GoogleComputeSnapshot(compute_base_resource.GoogleComputeBaseResource):
   """Class representing a Compute Engine Snapshot.
diff --git a/tools/cli.py b/tools/cli.py
index dca8c3e..32d717f 100644
--- a/tools/cli.py
+++ b/tools/cli.py
@@ -203,8 +203,9 @@ def Main() -> None:
                                 'instance will be copied.', None),
                 ('--disk_type', 'The SKU name for the disk to create. '
                                 'Can be Standard_LRS, Premium_LRS, '
-                                'StandardSSD_LRS, or UltraSSD_LRS. Default is '
-                                'Standard_LRS', 'Standard_LRS'),
+                                'StandardSSD_LRS, or UltraSSD_LRS. The default '
+                                'behavior is to use the same disk type as '
+                                'the source disk.', None),
                 ('--region', 'The region in which to create the disk copy. If '
                              'not provided, the disk copy will be created in '
                              'the "eastus" region.', 'eastus'),
@@ -283,7 +284,10 @@ def Main() -> None:
                 ('--disk_name', 'Name of the disk to copy. If none specified, '
                                 'then --instance_name must be specified and '
                                 'the boot disk of the instance will be copied.',
-                 None)
+                 None),
+                ('--disk_type', 'Type of disk. Can be pd-standard or pd-ssd. '
+                                'The default behavior is to use the same disk '
+                                'type as the source disk.', None)
             ])
   AddParser('gcp', gcp_subparsers, 'startvm', 'Start a forensic analysis VM.',
             args=[
@@ -291,7 +295,8 @@ def Main() -> None:
                  ''),
                 ('zone', 'Zone to create the instance in.', ''),
                 ('--disk_size', 'Size of disk in GB.', '50'),
-                ('--disk_type', 'Type of disk.', 'pd-ssd'),
+                ('--disk_type', 'Type of disk. Can be pd-standard or pd-ssd. '
+                                'The default value is pd-ssd.', 'pd-ssd'),
                 ('--cpu_cores', 'Instance CPU core count.', '4'),
                 ('--attach_disks', 'Comma separated list of disk names '
                                    'to attach.', None)
diff --git a/tools/gcp_cli.py b/tools/gcp_cli.py
index 97359f0..614c8ea 100644
--- a/tools/gcp_cli.py
+++ b/tools/gcp_cli.py
@@ -78,7 +78,8 @@ def CreateDiskCopy(args: 'argparse.Namespace') -> None:
                                   args.dst_project,
                                   args.zone,
                                   instance_name=args.instance_name,
-                                  disk_name=args.disk_name)
+                                  disk_name=args.disk_name,
+                                  disk_type=args.disk_type)
 
   logger.info('Disk copy completed.')
   logger.info('Name: {0:s}'.format(disk.name))


Test patch:
diff --git a/tests/providers/azure/azure_test.py b/tests/providers/azure/azure_test.py
index ce31b44..eae92dc 100644
--- a/tests/providers/azure/azure_test.py
+++ b/tests/providers/azure/azure_test.py
@@ -18,6 +18,8 @@ import typing
 import unittest
 import mock
 
+from azure.mgmt.compute.v2020_05_01 import models  # pylint: disable=import-error
+
 from libcloudforensics.providers.azure.internal import account, compute, common, monitoring  # pylint:disable=line-too-long
 from libcloudforensics.providers.azure import forensics
 
@@ -253,8 +255,7 @@ class TestAccount(unittest.TestCase):
     mock_create_disk.assert_called_with(
         FAKE_SNAPSHOT.resource_group_name,
         'fake_snapshot_name_f4c186ac_copy',
-        mock.ANY,
-        sku='Standard_LRS')
+        mock.ANY)
 
     # CreateDiskFromSnapshot(
     #     snapshot=FAKE_SNAPSHOT,
@@ -265,8 +266,7 @@ class TestAccount(unittest.TestCase):
     mock_create_disk.assert_called_with(
         FAKE_SNAPSHOT.resource_group_name,
         'new-forensics-disk',
-        mock.ANY,
-        sku='Standard_LRS')
+        mock.ANY)
 
     # CreateDiskFromSnapshot(
     #     snapshot=FAKE_SNAPSHOT, disk_name=None, disk_name_prefix='prefix')
@@ -275,18 +275,26 @@ class TestAccount(unittest.TestCase):
     mock_create_disk.assert_called_with(
         FAKE_SNAPSHOT.resource_group_name,
         'prefix_fake_snapshot_name_f4c186ac_copy',
-        mock.ANY,
-        sku='Standard_LRS')
+        mock.ANY)
 
     # CreateDiskFromSnapshot(
     #     snapshot=FAKE_SNAPSHOT, disk_type='StandardSSD_LRS')
+    expected_args = {
+        'location': 'fake-region',
+        'creation_data': {
+            'sourceResourceId': '/a/b/c/fake-resource-group/fake_snapshot_name',
+            'create_option': models.DiskCreateOption.copy
+        },
+        'sku': {
+            'name': 'StandardSSD_LRS'
+        }
+    }
     FAKE_ACCOUNT.CreateDiskFromSnapshot(
         FAKE_SNAPSHOT, disk_type='StandardSSD_LRS')
     mock_create_disk.assert_called_with(
         FAKE_SNAPSHOT.resource_group_name,
         'fake_snapshot_name_f4c186ac_copy',
-        mock.ANY,
-        sku='StandardSSD_LRS')
+        expected_args)
 
   @mock.patch('azure.storage.blob._container_client.ContainerClient.create_container')
   @mock.patch('azure.storage.blob._generated._azure_blob_storage.AzureBlobStorage.__init__')
@@ -330,8 +338,7 @@ class TestAccount(unittest.TestCase):
     mock_create_disk.assert_called_with(
         FAKE_SNAPSHOT.resource_group_name,
         'fake_snapshot_name_f4c186ac_copy',
-        mock.ANY,
-        sku='Standard_LRS')
+        mock.ANY)
 
   @mock.patch('sshpubkeys.SSHKey.parse')
   @mock.patch('libcloudforensics.scripts.utils.ReadStartupScript')
@@ -613,8 +620,9 @@ class TestMonitoring(unittest.TestCase):
 
 class TestForensics(unittest.TestCase):
   """Test Azure forensics file."""
-  # pylint: disable=line-too-long
+  # pylint: disable=line-too-long, too-many-arguments
 
+  @mock.patch('libcloudforensics.providers.azure.internal.compute.AZDisk.GetDiskType')
   @mock.patch('libcloudforensics.providers.azure.internal.account.AZAccount._GetOrCreateResourceGroup')
   @mock.patch('libcloudforensics.providers.azure.internal.common.GetCredentials')
   @mock.patch('libcloudforensics.providers.azure.internal.account.AZAccount.ListSubscriptionIDs')
@@ -634,7 +642,8 @@ class TestForensics(unittest.TestCase):
                           mock_snapshot_delete,
                           mock_list_subscription_ids,
                           mock_credentials,
-                          mock_resource_group):
+                          mock_resource_group,
+                          mock_disk_type):
     """Test that a disk copy is correctly created.
 
     CreateDiskCopy(zone, instance_name='fake-vm-name', region='fake-region').
@@ -649,6 +658,7 @@ class TestForensics(unittest.TestCase):
     mock_list_subscription_ids.return_value = ['fake-subscription-id']
     mock_credentials.return_value = ('fake-subscription-id', mock.Mock())
     mock_resource_group.return_value = 'fake-resource-group'
+    mock_disk_type.return_value = 'fake-disk-type'
 
     disk_copy = forensics.CreateDiskCopy(
         FAKE_ACCOUNT.default_resource_group_name,
@@ -661,6 +671,7 @@ class TestForensics(unittest.TestCase):
     self.assertIsInstance(disk_copy, compute.AZDisk)
     self.assertEqual('fake_snapshot_name_f4c186ac_copy', disk_copy.name)
 
+  @mock.patch('libcloudforensics.providers.azure.internal.compute.AZDisk.GetDiskType')
   @mock.patch('libcloudforensics.providers.azure.internal.account.AZAccount._GetOrCreateResourceGroup')
   @mock.patch('libcloudforensics.providers.azure.internal.common.GetCredentials')
   @mock.patch('libcloudforensics.providers.azure.internal.account.AZAccount.ListSubscriptionIDs')
@@ -680,7 +691,8 @@ class TestForensics(unittest.TestCase):
                           mock_snapshot_delete,
                           mock_list_subscription_ids,
                           mock_credentials,
-                          mock_resource_group):
+                          mock_resource_group,
+                          mock_disk_type):
     """Test that a disk copy is correctly created.
 
     CreateDiskCopy(zone, disk_name='fake-disk-name', region='fake-region').
@@ -694,6 +706,7 @@ class TestForensics(unittest.TestCase):
     mock_list_subscription_ids.return_value = ['fake-subscription-id']
     mock_credentials.return_value = ('fake-subscription-id', mock.Mock())
     mock_resource_group.return_value = 'fake-resource-group'
+    mock_disk_type.return_value = 'fake-disk-type'
 
     disk_copy = forensics.CreateDiskCopy(
         FAKE_ACCOUNT.default_resource_group_name,
@@ -706,6 +719,7 @@ class TestForensics(unittest.TestCase):
     self.assertIsInstance(disk_copy, compute.AZDisk)
     self.assertEqual('fake_snapshot_name_f4c186ac_copy', disk_copy.name)
 
+  @mock.patch('libcloudforensics.providers.azure.internal.compute.AZDisk.GetDiskType')
   @mock.patch('libcloudforensics.providers.azure.internal.account.AZAccount._GetOrCreateResourceGroup')
   @mock.patch('libcloudforensics.providers.azure.internal.common.GetCredentials')
   @mock.patch('libcloudforensics.providers.azure.internal.account.AZAccount.ListSubscriptionIDs')
@@ -717,7 +731,8 @@ class TestForensics(unittest.TestCase):
                           mock_list_disk,
                           mock_list_subscription_ids,
                           mock_credentials,
-                          mock_resource_group):
+                          mock_resource_group,
+                          mock_disk_type):
     """Test that a disk copy is correctly created.
 
     The first call should raise a RuntimeError in GetInstance as we are
@@ -728,6 +743,7 @@ class TestForensics(unittest.TestCase):
     mock_list_subscription_ids.return_value = ['fake-subscription-id']
     mock_credentials.return_value = ('fake-subscription-id', mock.Mock())
     mock_resource_group.return_value = 'fake-resource-group'
+    mock_disk_type.return_value = 'fake-disk-type'
 
     with self.assertRaises(RuntimeError) as error:
       forensics.CreateDiskCopy(
diff --git a/tests/providers/gcp/gcp_test.py b/tests/providers/gcp/gcp_test.py
index 8f10deb..460d0fd 100644
--- a/tests/providers/gcp/gcp_test.py
+++ b/tests/providers/gcp/gcp_test.py
@@ -853,6 +853,7 @@ class GCPTest(unittest.TestCase):
   # pylint: disable=line-too-long
 
   @typing.no_type_check
+  @mock.patch('libcloudforensics.providers.gcp.internal.compute.GoogleComputeDisk.GetDiskType')
   @mock.patch('libcloudforensics.providers.gcp.internal.common.GoogleCloudComputeClient.BlockOperation')
   @mock.patch('libcloudforensics.providers.gcp.internal.compute.GoogleComputeInstance.GetDisk')
   @mock.patch('libcloudforensics.providers.gcp.internal.compute.GoogleComputeInstance.GetBootDisk')
@@ -863,7 +864,8 @@ class GCPTest(unittest.TestCase):
                           mock_get_instance,
                           mock_get_boot_disk,
                           mock_get_disk,
-                          mock_block_operation):
+                          mock_block_operation,
+                          mock_disk_type):
     """Test that a disk from a remote project is duplicated and attached to
     an analysis project. """
     instances = mock_gce_api.return_value.instances.return_value.aggregatedList
@@ -871,6 +873,7 @@ class GCPTest(unittest.TestCase):
     mock_get_instance.return_value = FAKE_INSTANCE
     mock_get_boot_disk.return_value = FAKE_BOOT_DISK
     mock_block_operation.return_value = None
+    mock_disk_type.return_value = 'fake-disk-type'
 
     # create_disk_copy(
     #     src_proj,
@@ -890,6 +893,7 @@ class GCPTest(unittest.TestCase):
     self.assertTrue(new_disk.name.endswith('-copy'))
 
   @typing.no_type_check
+  @mock.patch('libcloudforensics.providers.gcp.internal.compute.GoogleComputeDisk.GetDiskType')
   @mock.patch('libcloudforensics.providers.gcp.internal.common.GoogleCloudComputeClient.BlockOperation')
   @mock.patch('libcloudforensics.providers.gcp.internal.compute.GoogleComputeInstance.GetBootDisk')
   @mock.patch('libcloudforensics.providers.gcp.internal.compute.GoogleCloudCompute.GetDisk')
@@ -898,13 +902,15 @@ class GCPTest(unittest.TestCase):
                           mock_gce_api,
                           mock_get_disk,
                           mock_get_boot_disk,
-                          mock_block_operation):
+                          mock_block_operation,
+                          mock_disk_type):
     """Test that a disk from a remote project is duplicated and attached to
     an analysis project. """
     instances = mock_gce_api.return_value.instances.return_value.aggregatedList
     instances.return_value.execute.return_value = MOCK_INSTANCES_AGGREGATED
     mock_get_disk.return_value = FAKE_DISK
     mock_block_operation.return_value = None
+    mock_disk_type.return_value = 'fake-disk-type'
 
     # create_disk_copy(
     #     src_proj,