Add local testing and fix function response

Author: davidgamez

.gitignore

@@ -77,4 +77,7 @@ coverage_reports
 tf.plan
 
 # CSV generation output files
-functions-python/**/*.csv
+functions-python/**/*.csv
+
+# Local emulators
+.cloudstorage

functions-python/batch_process_dataset/requirements_dev.txt

@@ -2,4 +2,5 @@ Faker
 pytest~=7.4.3
 urllib3-mock
 requests-mock
-python-dotenv~=1.0.0
+python-dotenv~=1.0.0
+gcp-storage-emulator

functions-python/batch_process_dataset/src/main.py

@@ -131,7 +131,11 @@ def download_content(self, temporary_file_path):
         return file_hash, is_zip
 
     def upload_file_to_storage(
-        self, source_file_path, dataset_stable_id, extracted_files_path
+        self,
+        source_file_path,
+        dataset_stable_id,
+        extracted_files_path,
+        public=True,
     ):
         """
         Uploads a file to the GCP bucket
@@ -146,7 +150,8 @@ def upload_file_to_storage(
             blob = bucket.blob(target_path)
             with open(source_file_path, "rb") as file:
                 blob.upload_from_file(file)
-            blob.make_public()
+            if public:
+                blob.make_public()
 
         base_path, _ = os.path.splitext(source_file_path)
         extracted_files: List[Gtfsfile] = []
@@ -162,7 +167,8 @@ def upload_file_to_storage(
                     f"{self.feed_stable_id}/{dataset_stable_id}/extracted/{file_name}"
                 )
                 file_blob.upload_from_filename(file_path)
-                file_blob.make_public()
+                if public:
+                    file_blob.make_public()
                 self.logger.info(
                     f"Uploaded extracted file {file_name} to {file_blob.public_url}"
                 )
@@ -175,7 +181,7 @@ def upload_file_to_storage(
                 )
         return blob, extracted_files
 
-    def upload_dataset(self) -> DatasetFile or None:
+    def upload_dataset(self, public=True) -> DatasetFile or None:
         """
         Uploads a dataset to a GCP bucket as <feed_stable_id>/latest.zip and
         <feed_stable_id>/<feed_stable_id>-<upload_datetime>.zip
@@ -185,9 +191,7 @@ def upload_dataset(self) -> DatasetFile or None:
         try:
             self.logger.info("Accessing URL %s", self.producer_url)
             temp_file_path = self.generate_temp_filename()
-            file_sha256_hash, is_zip, extracted_files_path = self.download_content(
-                temp_file_path
-            )
+            file_sha256_hash, is_zip = self.download_content(temp_file_path)
             if not is_zip:
                 self.logger.error(
                     f"[{self.feed_stable_id}] The downloaded file from {self.producer_url} is not a valid ZIP file."
@@ -202,9 +206,7 @@ def upload_dataset(self) -> DatasetFile or None:
                     f"[{self.feed_stable_id}] Dataset has changed (hash {self.latest_hash}"
                     f"-> {file_sha256_hash}). Uploading new version."
                 )
-                extracted_files_path = self.unzip_files(
-                    extracted_files_path, temp_file_path
-                )
+                extracted_files_path = self.unzip_files(temp_file_path)
                 self.logger.info(
                     f"Creating file {self.feed_stable_id}/latest.zip in bucket {self.bucket_name}"
                 )
@@ -220,7 +222,10 @@ def upload_dataset(self) -> DatasetFile or None:
                     f" in bucket {self.bucket_name}"
                 )
                 _, extracted_files = self.upload_file_to_storage(
-                    temp_file_path, dataset_stable_id, extracted_files_path
+                    temp_file_path,
+                    dataset_stable_id,
+                    extracted_files_path,
+                    public=public,
                 )
 
                 return DatasetFile(
@@ -242,7 +247,7 @@ def upload_dataset(self) -> DatasetFile or None:
                 os.remove(temp_file_path)
         return None
 
-    def unzip_files(self, extracted_files_path, temp_file_path):
+    def unzip_files(self, temp_file_path):
         extracted_files_path = os.path.join(temp_file_path.split(".")[0], "extracted")
         self.logger.info(f"Unzipping files to {extracted_files_path}")
         # Create the directory for extracted files if it does not exist

functions-python/batch_process_dataset/src/scripts/download_verifier.py

@@ -2,6 +2,12 @@
 import os
 
 from main import DatasetProcessor
+from gcp_storage_emulator.server import create_server
+
+HOST = "localhost"
+PORT = 9023
+BUCKET_NAME = "verifier"
+PRODUCER_URL = "https://example.com/dataset.zip"  # Replace with actual producer URL
 
 
 def verify_download_content(producer_url: str):
@@ -27,17 +33,52 @@ def verify_download_content(producer_url: str):
     )
     tempfile = processor.generate_temp_filename()
     logging.info(f"Temp filename: {tempfile}")
-    file_hash, is_zip, extracted_files_path = processor.download_content(tempfile)
-    are_files_extracted = os.path.exists(extracted_files_path)
+    file_hash, is_zip = processor.download_content(tempfile)
     logging.info(f"File hash: {file_hash}")
-    logging.info(
-        f"File path: {extracted_files_path} "
-        f"- { 'Files extracted' if are_files_extracted else 'Files not extracted'}"
+
+
+def verify_upload_dataset(producer_url: str):
+    """
+    Verifies the upload_dataset is able to upload the dataset to the GCP storage emulator.
+    This is useful to simulate the upload code locally and test issues related with user-agent and uploaded content.
+    This function also tests the DatasetProcessor class methods for generating a temporary filename
+    and uploading the dataset.
+    :param producer_url:
+    :return:
+    """
+    processor = DatasetProcessor(
+        producer_url=producer_url,
+        feed_id="feed_id",
+        feed_stable_id="feed_stable_id",
+        execution_id=None,
+        latest_hash="123",
+        bucket_name=BUCKET_NAME,
+        authentication_type=0,
+        api_key_parameter_name=None,
+        public_hosted_datasets_url=None,
     )
-    logging.info(f"Downloaded file path: {extracted_files_path}")
+    tempfile = processor.generate_temp_filename()
+    logging.info(f"Temp filename: {tempfile}")
+    dataset_file = processor.upload_dataset(public=False)
+    logging.info(f"Dataset File: {dataset_file}")
 
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)
-    # Example usage: Replace the placeholder URL with the actual producer URL before running.
-    verify_download_content(producer_url="https://example.com/your-dataset.zip")
+    # Replace with actual producer URL
+    try:
+        os.environ["STORAGE_EMULATOR_HOST"] = f"http://{HOST}:{PORT}"
+        server = create_server(
+            host=HOST, port=PORT, in_memory=False, default_bucket=BUCKET_NAME
+        )
+        server.start()
+
+        verify_download_content(producer_url=PRODUCER_URL)
+        logging.info("Download content verification completed successfully.")
+        verify_upload_dataset(producer_url=PRODUCER_URL)
+        verify_upload_dataset(producer_url=PRODUCER_URL)
+    except Exception as e:
+        logging.error(f"Error verifying download content: {e}")
+    finally:
+        server.stop()
+        logging.info("Verification completed.")

functions-python/batch_process_dataset/tests/test_batch_process_dataset_main.py

@@ -57,7 +57,7 @@ def test_upload_dataset_diff_hash(
         mock_blob.public_url = public_url
         mock_blob.path = public_url
         upload_file_to_storage.return_value = mock_blob, []
-        mock_download_url_content.return_value = file_hash, True, "path/file"
+        mock_download_url_content.return_value = file_hash, True
         mock_unzip_files.return_value = [mock_blob, mock_blob]
 
         processor = DatasetProcessor(
@@ -96,7 +96,7 @@ def test_upload_dataset_same_hash(
         mock_blob = MagicMock()
         mock_blob.public_url = public_url
         upload_file_to_storage.return_value = mock_blob
-        mock_download_url_content.return_value = file_hash, True, "path/file"
+        mock_download_url_content.return_value = file_hash, True
 
         processor = DatasetProcessor(
             public_url,
@@ -128,7 +128,7 @@ def test_upload_dataset_not_zip(
         mock_blob = MagicMock()
         mock_blob.public_url = public_url
         upload_file_to_storage.return_value = mock_blob
-        mock_download_url_content.return_value = file_hash, False, "path/file"
+        mock_download_url_content.return_value = file_hash, False
 
         processor = DatasetProcessor(
             public_url,

functions-python/helpers/logger.py

@@ -42,6 +42,7 @@ def filter(self, record):
 
 
 lock = threading.Lock()
+lock_logger = threading.Lock()
 _logging_initialized = False
 
 
@@ -73,14 +74,15 @@ def get_logger(name: str, stable_id: str = None):
     If stable_id is provided, the StableIdFilter is added.
     This method can be called multiple times for the same logger name without creating a side effect.
     """
-    # Create the logger with the provided name to avoid retuning the same logger instance
-    logger = (
-        logging.getLogger(name)
-        if not stable_id
-        else logging.getLogger(f"{name}_{stable_id}")
-    )
-    if stable_id and not any(
-        isinstance(log_filter, StableIdFilter) for log_filter in logger.filters
-    ):
-        logger.addFilter(StableIdFilter(stable_id))
-    return logger
+    with lock_logger:
+        # Create the logger with the provided name to avoid retuning the same logger instance
+        logger = (
+            logging.getLogger(f"{name}_{stable_id}")
+            if stable_id
+            else logging.getLogger(name)
+        )
+        if stable_id and not any(
+            isinstance(log_filter, StableIdFilter) for log_filter in logger.filters
+        ):
+            logger.addFilter(StableIdFilter(stable_id))
+        return logger