fix: unit test

MobilityData · Jul 25, 2024 · aab1f6e · aab1f6e
1 parent 2f3b06b
commit aab1f6e
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/functions-python/extract_bb/src/main.py b/functions-python/extract_bb/src/main.py
@@ -247,6 +247,7 @@ def extract_bounding_box_batch(_):
 
     # Get latest GTFS dataset with no bounding boxes
     session = None
+    execution_id = str(uuid.uuid4())
     datasets_data = []
     try:
         session = start_db_session(os.getenv("FEEDS_DATABASE_URL"))
@@ -261,6 +262,7 @@ def extract_bounding_box_batch(_):
                 "stable_id": dataset.feed_id,
                 "dataset_id": dataset.stable_id,
                 "url": dataset.hosted_url,
+                "execution_id": execution_id,
             }
             datasets_data.append(data)
             logging.info(f"Dataset {dataset.stable_id} added to the batch.")

diff --git a/functions-python/extract_bb/tests/test_extract_bb.py b/functions-python/extract_bb/tests/test_extract_bb.py
@@ -71,7 +71,8 @@ def test_get_gtfs_feed_bounds(self, mock_gtfs_kit):
             self.assertEqual(bounds[i], expected_bounds[i])
 
     @patch("extract_bb.src.main.Logger")
-    def test_extract_bb_exception(self, _):
+    @patch("extract_bb.src.main.DatasetTraceService")
+    def test_extract_bb_exception(self, _, __):
         # Data with missing url
         data = {"stable_id": faker.pystr(), "dataset_id": faker.pystr()}
         message_data = base64.b64encode(json.dumps(data).encode("utf-8")).decode(
@@ -111,6 +112,7 @@ def test_extract_bb_exception(self, _):
         os.environ,
         {
             "FEEDS_DATABASE_URL": default_db_url,
+            "GOOGLE_APPLICATION_CREDENTIALS": "dummy-credentials.json"
         },
     )
     @patch("extract_bb.src.main.get_gtfs_feed_bounds")
@@ -153,14 +155,16 @@ def test_extract_bb(
         {
             "FEEDS_DATABASE_URL": default_db_url,
             "MAXIMUM_EXECUTIONS": "1",
+            "GOOGLE_APPLICATION_CREDENTIALS": "dummy-credentials.json"
         },
     )
     @patch("extract_bb.src.main.get_gtfs_feed_bounds")
     @patch("extract_bb.src.main.update_dataset_bounding_box")
     @patch("extract_bb.src.main.DatasetTraceService.get_by_execution_and_stable_ids")
     @patch("extract_bb.src.main.Logger")
+    @patch("google.cloud.datastore.Client")
     def test_extract_bb_max_executions(
-        self, _, mock_dataset_trace, update_bb_mock, get_gtfs_feed_bounds_mock
+        self, _, __, mock_dataset_trace, update_bb_mock, get_gtfs_feed_bounds_mock
     ):
         get_gtfs_feed_bounds_mock.return_value = np.array(
             [faker.longitude(), faker.latitude(), faker.longitude(), faker.latitude()]
@@ -193,6 +197,7 @@ def test_extract_bb_max_executions(
         os.environ,
         {
             "FEEDS_DATABASE_URL": default_db_url,
+            "GOOGLE_APPLICATION_CREDENTIALS": "dummy-credentials.json"
         },
     )
     @patch("extract_bb.src.main.get_gtfs_feed_bounds")
@@ -228,6 +233,7 @@ def test_extract_bb_cloud_event(
         os.environ,
         {
             "FEEDS_DATABASE_URL": default_db_url,
+            "GOOGLE_APPLICATION_CREDENTIALS": "dummy-credentials.json"
         },
     )
     @patch("extract_bb.src.main.get_gtfs_feed_bounds")
@@ -255,6 +261,7 @@ def test_extract_bb_cloud_event_error(
         os.environ,
         {
             "FEEDS_DATABASE_URL": default_db_url,
+            "GOOGLE_APPLICATION_CREDENTIALS": "dummy-credentials.json"
         },
     )
     @patch("extract_bb.src.main.get_gtfs_feed_bounds")
@@ -296,13 +303,15 @@ def test_extract_bb_exception_2(self, _, update_bb_mock, get_gtfs_feed_bounds_mo
             "FEEDS_DATABASE_URL": default_db_url,
             "PUBSUB_TOPIC_NAME": "test-topic",
             "PROJECT_ID": "test-project",
+            "GOOGLE_APPLICATION_CREDENTIALS": "dummy-credentials.json"
         },
     )
     @patch("extract_bb.src.main.start_db_session")
     @patch("extract_bb.src.main.pubsub_v1.PublisherClient")
     @patch("extract_bb.src.main.Logger")
+    @patch("uuid.uuid4")
     def test_extract_bounding_box_batch(
-        self, logger_mock, publisher_client_mock, start_db_session_mock
+        self, uuid_mock, logger_mock, publisher_client_mock, start_db_session_mock
     ):
         # Mock the database session and query
         mock_session = MagicMock()
@@ -324,6 +333,7 @@ def test_extract_bounding_box_batch(
             mock_dataset1,
             mock_dataset2,
         ]
+        uuid_mock.return_value = "batch-uuid"
         start_db_session_mock.return_value = mock_session
 
         # Mock the Pub/Sub client
@@ -345,6 +355,7 @@ def test_extract_bounding_box_batch(
                     "stable_id": "1",
                     "dataset_id": "stable_1",
                     "url": "http://example.com/1",
+                    "execution_id": "batch-uuid",
                 }
             ).encode("utf-8"),
         )
@@ -355,6 +366,7 @@ def test_extract_bounding_box_batch(
                     "stable_id": "2",
                     "dataset_id": "stable_2",
                     "url": "http://example.com/2",
+                    "execution_id": "batch-uuid",
                 }
             ).encode("utf-8"),
         )
@@ -364,6 +376,7 @@ def test_extract_bounding_box_batch(
         os.environ,
         {
             "FEEDS_DATABASE_URL": default_db_url,
+            "GOOGLE_APPLICATION_CREDENTIALS": "dummy-credentials.json"
         },
     )
     @patch("extract_bb.src.main.Logger")
@@ -379,6 +392,7 @@ def test_extract_bounding_box_batch_no_topic_name(self, logger_mock):
             "FEEDS_DATABASE_URL": default_db_url,
             "PUBSUB_TOPIC_NAME": "test-topic",
             "PROJECT_ID": "test-project",
+            "GOOGLE_APPLICATION_CREDENTIALS": "dummy-credentials.json"
         },
     )
     @patch("extract_bb.src.main.start_db_session")