From 2ebd33c7e0feb6ac6413fa3e147a0539d1e00631 Mon Sep 17 00:00:00 2001 From: David Gamez <1192523+davidgamez@users.noreply.github.com> Date: Wed, 21 Aug 2024 09:29:40 -0400 Subject: [PATCH] fix: fix invalid feeds credentials variable name (#704) --- .../batch_process_dataset/src/main.py | 24 +++- .../tests/test_batch_process_dataset_main.py | 103 ++++++++++++++++++ 2 files changed, 124 insertions(+), 3 deletions(-) diff --git a/functions-python/batch_process_dataset/src/main.py b/functions-python/batch_process_dataset/src/main.py index 15448347e..68ac63f00 100644 --- a/functions-python/batch_process_dataset/src/main.py +++ b/functions-python/batch_process_dataset/src/main.py @@ -25,7 +25,6 @@ from typing import Optional import functions_framework -import ast from cloudevents.http import CloudEvent from google.cloud import storage from sqlalchemy import func @@ -76,13 +75,32 @@ def __init__( self.authentication_type = authentication_type self.api_key_parameter_name = api_key_parameter_name self.date = datetime.now().strftime("%Y%m%d%H%M") - feeds_credentials = ast.literal_eval(os.getenv("FEED_CREDENTIALS", "{}")) - self.feed_credentials = feeds_credentials.get(self.feed_stable_id, None) + if self.authentication_type != 0: + logging.info(f"Getting feed credentials for feed {self.feed_stable_id}") + self.feed_credentials = self.get_feed_credentials(self.feed_stable_id) + if self.feed_credentials is None: + raise Exception( + f"Error getting feed credentials for feed {self.feed_stable_id}" + ) + else: + self.feed_credentials = None self.public_hosted_datasets_url = public_hosted_datasets_url self.init_status = None self.init_status_additional_data = None + @staticmethod + def get_feed_credentials(feed_stable_id) -> str | None: + """ + Gets the feed credentials from the environment variable + """ + try: + feeds_credentials = json.loads(os.getenv("FEEDS_CREDENTIALS", "{}")) + return feeds_credentials.get(feed_stable_id, None) + except Exception as e: + logging.error(f"Error getting feed credentials: {e}") + return None + @staticmethod def create_dataset_stable_id(feed_stable_id, timestamp): """ diff --git a/functions-python/batch_process_dataset/tests/test_batch_process_dataset_main.py b/functions-python/batch_process_dataset/tests/test_batch_process_dataset_main.py index 830378f8b..13f6e6fc1 100644 --- a/functions-python/batch_process_dataset/tests/test_batch_process_dataset_main.py +++ b/functions-python/batch_process_dataset/tests/test_batch_process_dataset_main.py @@ -211,6 +211,9 @@ def test_upload_file_to_storage(self): # Assert that the file was opened in binary read mode mock_file.assert_called_once_with(source_file_path, "rb") + @patch.dict( + os.environ, {"FEEDS_CREDENTIALS": '{"test_stable_id": "test_credentials"}'} + ) def test_process(self): session = get_testing_session() feeds = session.query(Gtfsfeed).all() @@ -252,6 +255,77 @@ def test_process(self): self.assertEqual(result.file_sha256_hash, new_hash) processor.upload_dataset.assert_called_once() + @patch.dict( + os.environ, + {"FEEDS_CREDENTIALS": '{"not_what_u_r_looking_4": "test_credentials"}'}, + ) + def test_fails_authenticated_feed_not_creds(self): + session = get_testing_session() + feeds = session.query(Gtfsfeed).all() + feed_id = feeds[0].id + + producer_url = "https://testproducer.com/data" + feed_stable_id = "test_stable_id" + execution_id = "test_execution_id" + latest_hash = "old_hash" + bucket_name = "test-bucket" + authentication_type = 1 + api_key_parameter_name = "test_api_key" + + with self.assertRaises(Exception) as context: + DatasetProcessor( + producer_url, + feed_id, + feed_stable_id, + execution_id, + latest_hash, + bucket_name, + authentication_type, + api_key_parameter_name, + test_hosted_public_url, + ) + self.assertEqual( + str(context.exception), + "Error getting feed credentials for feed test_stable_id", + ) + + @patch.dict( + os.environ, + {"FEEDS_CREDENTIALS": "not a JSON string"}, + ) + def test_fails_authenticated_feed_creds_invalid(self): + session = get_testing_session() + feeds = session.query(Gtfsfeed).all() + feed_id = feeds[0].id + + producer_url = "https://testproducer.com/data" + feed_stable_id = "test_stable_id" + execution_id = "test_execution_id" + latest_hash = "old_hash" + bucket_name = "test-bucket" + authentication_type = 1 + api_key_parameter_name = "test_api_key" + + with self.assertRaises(Exception) as context: + DatasetProcessor( + producer_url, + feed_id, + feed_stable_id, + execution_id, + latest_hash, + bucket_name, + authentication_type, + api_key_parameter_name, + test_hosted_public_url, + ) + self.assertEqual( + str(context.exception), + "Error getting feed credentials for feed test_stable_id", + ) + + @patch.dict( + os.environ, {"FEEDS_CREDENTIALS": '{"test_stable_id": "test_credentials"}'} + ) def test_process_no_change(self): feed_id = "test" producer_url = "https://testproducer.com/data" @@ -344,3 +418,32 @@ def test_process_dataset_exception( assert False except AttributeError: assert True + + @patch("batch_process_dataset.src.main.Logger") + @patch("batch_process_dataset.src.main.DatasetTraceService") + def test_process_dataset_missing_stable_id(self, mock_dataset_trace, _): + db_url = os.getenv("TEST_FEEDS_DATABASE_URL", default=default_db_url) + os.environ["FEEDS_DATABASE_URL"] = db_url + + # Mock data for the CloudEvent + mock_data = { + "execution_id": "test_execution_id", + "producer_url": "https://testproducer.com/data", + "feed_stable_id": "", + "feed_id": "test_feed_id", + "dataset_id": "test_dataset_id", + "dataset_hash": "test_dataset_hash", + "authentication_type": 0, + "api_key_parameter_name": None, + } + + cloud_event = create_cloud_event(mock_data) + + mock_dataset_trace.save.return_value = None + mock_dataset_trace.get_by_execution_and_stable_ids.return_value = 0 + # Call the function + result = process_dataset(cloud_event) + assert ( + result + == "Function completed with errors, missing stable= or execution_id=test_execution_id" + )