fixed broken tests

qcdyx · qcdyx · commit 18d92abe75ff · 2025-08-05T20:30:21.000-04:00
diff --git a/api/src/shared/common/gcp_utils.py b/api/src/shared/common/gcp_utils.py
@@ -0,0 +1,100 @@
+import logging
+import os
+from google.cloud import tasks_v2
+from google.protobuf.timestamp_pb2 import Timestamp
+
+
+def create_refresh_materialized_view_task():
+    """
+    Asynchronously refresh a materialized view.
+    Ensures deduplication by generating a unique task name.
+
+    Returns:
+        dict: Response message and status code.
+    """
+    from google.protobuf import timestamp_pb2
+    from datetime import datetime, timedelta
+
+    try:
+        logging.info("Creating materialized view refresh task.")
+        now = datetime.now()
+
+        # BOUNCE WINDOW: next :00 or :30
+        minute = now.minute
+        if minute < 30:
+            bucket_time = now.replace(minute=30, second=0, microsecond=0)
+        else:
+            bucket_time = now.replace(minute=0, second=0, microsecond=0) + timedelta(hours=1)
+
+        timestamp_str = bucket_time.strftime("%Y-%m-%d-%H-%M")
+        task_name = f"refresh-materialized-view-{timestamp_str}"
+
+        # Convert to protobuf timestamp
+        proto_time = timestamp_pb2.Timestamp()
+        proto_time.FromDatetime(bucket_time)
+
+        # Cloud Tasks setup
+        client = tasks_v2.CloudTasksClient()
+        project = os.getenv("PROJECT_ID")
+        location = os.getenv("LOCATION")
+        queue = os.getenv("MATERIALIZED_VIEW_QUEUE")
+        url = (
+            f"https://{os.getenv('GCP_REGION')}-"
+            f"{os.getenv('PROJECT_ID')}.cloudfunctions.net/"
+            f"tasks-executor-{os.getenv('ENVIRONMENT_NAME')}"
+        )
+
+        task_name = client.task_path(project, location, queue, task_name)
+
+        # Enqueue the task
+        try:
+            create_http_task_with_name(
+                client=client,
+                body=b"",
+                url=url,
+                project_id=project,
+                gcp_region=location,
+                queue_name=queue,
+                task_name=task_name,
+                task_time=proto_time,
+                http_method=tasks_v2.HttpMethod.GET,
+            )
+            logging.info(f"Scheduled refresh materialized view task for {timestamp_str}")
+            return {"message": f"Refresh task for {timestamp_str} scheduled."}, 200
+        except Exception as e:
+            if "ALREADY_EXISTS" in str(e):
+                logging.info(f"Task already exists for {timestamp_str}, skipping.")
+
+    except Exception as error:
+        error_msg = f"Error enqueuing task: {error}"
+        logging.error(error_msg)
+        return {"error": error_msg}, 500
+
+
+def create_http_task_with_name(
+    client: "tasks_v2.CloudTasksClient",
+    body: bytes,
+    url: str,
+    project_id: str,
+    gcp_region: str,
+    queue_name: str,
+    task_name: str,
+    task_time: Timestamp,
+    http_method: "tasks_v2.HttpMethod",
+):
+    """Creates a GCP Cloud Task."""
+
+    token = tasks_v2.OidcToken(service_account_email=os.getenv("SERVICE_ACCOUNT_EMAIL"))
+
+    task = tasks_v2.Task(
+        name=task_name,
+        schedule_time=task_time,
+        http_request=tasks_v2.HttpRequest(
+            url=url,
+            http_method=http_method,
+            oidc_token=token,
+            body=body,
+            headers={"Content-Type": "application/json"},
+        ),
+    )
+    client.create_task(parent=client.queue_path(project_id, gcp_region, queue_name), task=task)
diff --git a/api/src/shared/database/database.py b/api/src/shared/database/database.py
@@ -21,6 +21,7 @@
 from sqlalchemy.orm import sessionmaker
 import logging
 
+
 from shared.common.logging_utils import get_env_logging_level
 
 
@@ -116,78 +117,6 @@ def mapper_configure_listener(mapper, class_):
 event.listen(mapper, "mapper_configured", mapper_configure_listener)
 
 
-def create_refresh_materialized_view_task():
-    """
-    Asynchronously refresh a materialized view.
-    Ensures deduplication by generating a unique task name.
-
-    Returns:
-        dict: Response message and status code.
-    """
-    from google.protobuf import timestamp_pb2
-    from datetime import datetime, timedelta
-    from google.cloud import tasks_v2
-
-    try:
-        logging.info("Creating materialized view refresh task.")
-        now = datetime.now()
-
-        # BOUNCE WINDOW: next :00 or :30
-        minute = now.minute
-        if minute < 30:
-            bucket_time = now.replace(minute=30, second=0, microsecond=0)
-        else:
-            bucket_time = now.replace(minute=0, second=0, microsecond=0) + timedelta(hours=1)
-
-        timestamp_str = bucket_time.strftime("%Y-%m-%d-%H-%M")
-        task_name = f"refresh-materialized-view-{timestamp_str}"
-
-        # Convert to protobuf timestamp
-        proto_time = timestamp_pb2.Timestamp()
-        proto_time.FromDatetime(bucket_time)
-
-        # Cloud Tasks setup
-        client = tasks_v2.CloudTasksClient()
-        project = os.getenv("PROJECT_ID")
-        location = os.getenv("LOCATION")
-        queue = os.getenv("MATERIALIZED_VIEW_QUEUE")
-        url = (
-            f"https://{os.getenv('GCP_REGION')}-"
-            f"{os.getenv('PROJECT_ID')}.cloudfunctions.net/"
-            f"tasks-executor-{os.getenv('ENVIRONMENT_NAME')}"
-        )
-
-        task_name = client.task_path(project, location, queue, task_name)
-
-        # Convert to protobuf timestamp
-        proto_time = timestamp_pb2.Timestamp()
-        proto_time.FromDatetime(bucket_time)
-
-        # Enqueue the task
-        try:
-            create_http_task_with_name(
-                client=client,
-                body=b"",
-                url=url,
-                project_id=project,
-                gcp_region=location,
-                queue_name=queue,
-                task_name=task_name,
-                task_time=proto_time,
-                http_method=tasks_v2.HttpMethod.GET,
-            )
-            logging.info(f"Scheduled refresh materialized view task for {timestamp_str}")
-            return {"message": f"Refresh task for {timestamp_str} scheduled."}, 200
-        except Exception as e:
-            if "ALREADY_EXISTS" in str(e):
-                logging.info(f"Task already exists for {timestamp_str}, skipping.")
-
-    except Exception as error:
-        error_msg = f"Error enqueuing task: {error}"
-        logging.error(error_msg)
-        return {"error": error_msg}, 500
-
-
 def refresh_materialized_view(session: "Session", view_name: str) -> bool:
     """
     Refresh Materialized view by name.
@@ -201,47 +130,6 @@ def refresh_materialized_view(session: "Session", view_name: str) -> bool:
         return False
 
 
-def create_http_task_with_name(
-    client,  # type: tasks_v2.CloudTasksClient
-    body: bytes,
-    url: str,
-    project_id: str,
-    gcp_region: str,
-    queue_name: str,
-    task_name: str,
-    task_time: str,
-    http_method,
-):
-    from google.cloud import tasks_v2
-
-    """Creates a GCP Cloud Task."""
-
-    token = (tasks_v2.OidcToken(service_account_email=os.getenv("SERVICE_ACCOUNT_EMAIL")),)
-
-    task = tasks_v2.Task(
-        http_request=tasks_v2.HttpRequest(
-            url=url,
-            http_method=http_method,
-            oidc_token=token,
-            body=body,
-            headers={"Content-Type": "application/json"},
-        )
-    )
-    task = {
-        "name": task_name,
-        "http_request": {
-            "http_method": http_method,
-            "url": url,
-            "headers": {
-                "Content-Type": "application/json",
-                "Authorization": f"Bearer {token}",
-            },
-        },
-        "schedule_time": task_time,
-    }
-    client.create_task(parent=client.queue_path(project_id, gcp_region, queue_name), task=task)
-
-
 def with_db_session(func=None, db_url: str | None = None):
     """
     Decorator to handle the session management for the decorated function.
diff --git a/functions-python/backfill_dataset_service_date_range/src/main.py b/functions-python/backfill_dataset_service_date_range/src/main.py
@@ -2,12 +2,10 @@
 import os
 import functions_framework
 
+from shared.common.gcp_utils import create_refresh_materialized_view_task
 from shared.helpers.logger import init_logger
 
-from shared.database.database import (
-    with_db_session,
-    create_refresh_materialized_view_task,
-)
+from shared.database.database import with_db_session
 
 from sqlalchemy.orm import joinedload, Session
 from sqlalchemy import or_, func
diff --git a/functions-python/batch_process_dataset/src/main.py b/functions-python/batch_process_dataset/src/main.py
@@ -29,13 +29,11 @@
 from google.cloud import storage
 from sqlalchemy import func
 
+from shared.common.gcp_utils import create_refresh_materialized_view_task
 from shared.database_gen.sqlacodegen_models import Gtfsdataset, Gtfsfile
 
 from shared.dataset_service.main import DatasetTraceService, DatasetTrace, Status
-from shared.database.database import (
-    with_db_session,
-    create_refresh_materialized_view_task,
-)
+from shared.database.database import with_db_session
 import logging
 
 from shared.helpers.logger import init_logger, get_logger
diff --git a/functions-python/helpers/feed_status.py b/functions-python/helpers/feed_status.py
@@ -6,9 +6,7 @@
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
-from shared.database.database import (
-    create_refresh_materialized_view_task,
-)
+from shared.common.gcp_utils import create_refresh_materialized_view_task
 
 
 #  query to update the status of the feeds based on the service date range of the latest dataset
diff --git a/functions-python/helpers/utils.py b/functions-python/helpers/utils.py
@@ -151,7 +151,12 @@ def create_http_task(
     gcp_region: str,
     queue_name: str,
 ) -> None:
-    from shared.database.database import create_http_task_with_name
+    from shared.common.gcp_utils import create_http_task_with_name
+    from google.cloud import tasks_v2
+    from google.protobuf import timestamp_pb2
+
+    proto_time = timestamp_pb2.Timestamp()
+    proto_time.GetCurrentTime()
 
     create_http_task_with_name(
         client=client,
@@ -160,4 +165,7 @@ def create_http_task(
         project_id=project_id,
         gcp_region=gcp_region,
         queue_name=queue_name,
+        task_name="task_name",
+        task_time=proto_time,
+        http_method=tasks_v2.HttpMethod.POST,
     )
diff --git a/functions-python/process_validation_report/requirements.txt b/functions-python/process_validation_report/requirements.txt
@@ -16,6 +16,7 @@ geoalchemy2==0.14.7
 
 # Google specific packages for this function
 cloudevents~=1.10.1
+google-cloud-tasks
 
 # Configuration
 python-dotenv==1.0.0
diff --git a/functions-python/reverse_geolocation/src/reverse_geolocation_processor.py b/functions-python/reverse_geolocation/src/reverse_geolocation_processor.py
@@ -23,10 +23,8 @@
     geopolygons_as_string,
 )
 from parse_request import parse_request_parameters
-from shared.database.database import (
-    with_db_session,
-    create_refresh_materialized_view_task,
-)
+from shared.common.gcp_utils import create_refresh_materialized_view_task
+from shared.database.database import with_db_session
 
 from shared.database_gen.sqlacodegen_models import (
     Geopolygon,
