Add snapshot property to Datasource class

ocp_resources/data_source.py

@@ -3,6 +3,7 @@
 from ocp_resources.constants import TIMEOUT_4MINUTES
 from ocp_resources.persistent_volume_claim import PersistentVolumeClaim
 from ocp_resources.resource import NamespacedResource
+from ocp_resources.volume_snapshot import VolumeSnapshot
 
 
 class DataSource(NamespacedResource):
@@ -39,19 +40,29 @@ def to_dict(self):
                 },
             })
 
-    @property
-    def pvc(self):
-        data_source_pvc = self.instance.spec.source.pvc
-        pvc_name = data_source_pvc.name
-        pvc_namespace = data_source_pvc.namespace
+    def _get_boot_source(self, boot_source_type):
+        boot_source = self.instance.spec.source.get(boot_source_type)
+        if not boot_source:
+            return None
+        boot_source_name = boot_source.name
+        boot_source_namespace = boot_source.namespace
         try:
-            return PersistentVolumeClaim(
+            boot_source_object = PersistentVolumeClaim if boot_source_type == "pvc" else VolumeSnapshot
+            return boot_source_object(
                 client=self.client,
-                name=pvc_name,
-                namespace=pvc_namespace,
+                name=boot_source_name,
+                namespace=boot_source_namespace,
             )
         except ResourceNotFoundError:
             self.logger.warning(
-                f"dataSource {self.name} is pointing to a non-existing PVC, name:"
-                f" {pvc_name}, namespace: {pvc_namespace}"
+                f"dataSource {self.name} is pointing to a non-existing {boot_source_type}, name:"
+                f" {boot_source_name}, namespace: {boot_source_namespace}"
             )
+
+    @property
+    def pvc(self):
+        return self._get_boot_source(boot_source_type="pvc")
+
+    @property
+    def snapshot(self):
+        return self._get_boot_source(boot_source_type="snapshot")