Demo Unit test for CloudStorageAPIClient

pyproject.toml

@@ -7,7 +7,7 @@ license = "MPL-2.0"
 readme = "README.md"
 
 [tool.poetry.dependencies]
-python = "^3.10"
+python = ">=3.10,<3.13"
 google-api-core = "^2.19.0"
 google-cloud-storage = "^2.16.0"
 wandb = "^0.16.6"
@@ -18,6 +18,7 @@ ruff = "^0.4.3"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.2.0"
+cloud-storage-mocker = "^0.3.4"
 
 [tool.pytest.ini_options]
 markers = [

src/mozmlops/cloud_storage_api_client.py

@@ -40,7 +40,7 @@ def store(self, data: bytes, storage_path: str) -> str:
         client = storage.Client(project=self.gcs_project_name)
 
         # Raises an exception if the bucket name cannot be found
-        bucket = client.get_bucket(self.gcs_bucket_name)
+        bucket = client.bucket(self.gcs_bucket_name)
 
         blob = bucket.blob(storage_path)
 
@@ -74,7 +74,7 @@ def fetch(self, remote_path: str, local_path: str) -> str:
         from google.cloud import storage
 
         client = storage.Client(project=self.gcs_project_name)
-        bucket = client.get_bucket(self.gcs_bucket_name)
+        bucket = client.bucket(self.gcs_bucket_name)
 
         blob = bucket.blob(remote_path)
 
@@ -83,6 +83,7 @@ def fetch(self, remote_path: str, local_path: str) -> str:
         p.parent.mkdir(parents=True, exist_ok=True)
 
         blob.download_to_filename(local_path)
+        return local_path
 
     def __delete(self, remote_path: str) -> str:
         """

tests/__init__.py

@@ -0,0 +1,6 @@
+import sys
+import os
+
+# To avoid ModuleNotFound errors in tests while attempting to import test subjects.
+# To be investigated here: https://mozilla-hub.atlassian.net/browse/DENG-3667
+sys.path.append(f"{os.getcwd()}/src")

tests/integration/__init__.py

@@ -1,6 +0,0 @@
-import sys
-import os
-
-# To avoid ModuleNotFound errors in tests while attempting to import test subjects.
-# To be investigated here: https://mozilla-hub.atlassian.net/browse/DENG-3667
-sys.path.append(f"{os.getcwd()}/src")

tests/unit/test_cloud_storage_api_client.py

@@ -0,0 +1,90 @@
+import io
+import pathlib
+from datetime import datetime
+
+import google.cloud.storage  # type: ignore[import]
+from cloud_storage_mocker import Mount
+from cloud_storage_mocker import patch as gcs_patch
+from mozmlops.cloud_storage_api_client import CloudStorageAPIClient
+
+
+def test_store__stores_file_on_gcs(tmp_path: pathlib.Path) -> None:
+    """
+    This is a unit test that checks whether we called the GCS API
+    according to the expectations of this GCS mocking library:
+
+    https://github.com/odashi/cloud-storage-mocker
+
+    The library does not support mocks for all GCS operations,
+    and it is not maintained by the team that builds the GCS API.
+
+    If this test fails and it's unclear why,
+    run the integration tests to check our integration behavior for real
+    to see if something is wrong:
+
+    pytest -m integration
+    """
+
+    # Given the following mocked bucket
+    with gcs_patch(
+        [
+            Mount("testbucket", tmp_path / "src", readable=True, writable=True),
+        ],
+    ):
+        # When our API Client stores data in a file on GCS:
+        storage_client = CloudStorageAPIClient(project_name="testproject", bucket_name="testbucket")
+
+        string_to_store = "Ada Lovelace"
+        timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+        filename_to_store_it_at = f"first_computer_programmer_{timestamp}.txt"
+        encoded_string = string_to_store.encode(encoding='utf-8')
+
+        filepath = storage_client.store(data=encoded_string, storage_path=filename_to_store_it_at)
+        assert filepath == filename_to_store_it_at
+
+        # Then we can do the steps to download the file from mock GCS:
+        mock_client = google.cloud.storage.Client()
+        blob = mock_client.bucket("testbucket").blob(filename_to_store_it_at)
+        assert blob.download_as_text() == string_to_store
+
+def test_fetch__gets_file_off_gcs(tmp_path: pathlib.Path) -> None:
+    """
+    This is a unit test that checks whether we called the GCS API
+    according to the expectations of this GCS mocking library:
+
+    https://github.com/odashi/cloud-storage-mocker
+
+    The library does not support mocks for all GCS operations,
+    and it is not maintained by the team that builds the GCS API.
+
+    If this test fails and it's unclear why,
+    run the integration tests to check our integration behavior for real
+    to see if something is wrong:
+
+    pytest -m integration
+    """
+
+    # Given the following mocked bucket
+    with gcs_patch(
+        [
+            Mount("testbucket", tmp_path / "src", readable=True, writable=True),
+        ],
+    ):
+        # When we do the steps to upload a file to mock GCS:
+        string_to_store = "Ada Lovelace"
+        timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+        filename_to_store_it_at = f"first_computer_programmer_{timestamp}.txt"
+        encoded_string = string_to_store.encode(encoding='utf-8')
+
+        mock_client = google.cloud.storage.Client()
+        blob = mock_client.bucket("testbucket").blob(filename_to_store_it_at)
+        with io.BytesIO(encoded_string) as data:
+            blob.upload_from_file(data)
+
+        # Then Our API client is able to fetch them:
+        storage_client = CloudStorageAPIClient(project_name="testproject", bucket_name="testbucket")
+        storage_client.fetch(remote_path=filename_to_store_it_at, local_path=tmp_path / filename_to_store_it_at)
+        assert (tmp_path / filename_to_store_it_at).read_text() == string_to_store
+
+
+