From ee6258572d35d845e2e44e4c57a5205c76eae7a9 Mon Sep 17 00:00:00 2001 From: Joshua Greben Date: Wed, 15 May 2024 14:33:19 -0700 Subject: [PATCH] Function to return success early for transmissions when not production --- .../data_exports/transmission_tasks.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/libsys_airflow/plugins/data_exports/transmission_tasks.py b/libsys_airflow/plugins/data_exports/transmission_tasks.py index dd1f6523..7fb2167c 100644 --- a/libsys_airflow/plugins/data_exports/transmission_tasks.py +++ b/libsys_airflow/plugins/data_exports/transmission_tasks.py @@ -4,6 +4,7 @@ import httpx from airflow.decorators import task +from airflow.models import Variable from airflow.models.connection import Connection from airflow.providers.ftp.hooks.ftp import FTPHook @@ -41,6 +42,8 @@ def gather_files_task(**kwargs) -> dict: @task(multiple_outputs=True) def transmit_data_http_task(gather_files, **kwargs) -> dict: + if not is_production(): + return return_success_test_instance(gather_files) """ Transmit the data via http Returns lists of files successfully transmitted and failures @@ -78,6 +81,8 @@ def transmit_data_http_task(gather_files, **kwargs) -> dict: @task(multiple_outputs=True) def transmit_data_ftp_task(conn_id, local_files) -> dict: + if not is_production(): + return return_success_test_instance(local_files) """ Transmit the data via ftp Returns lists of files successfully transmitted and failures @@ -128,3 +133,17 @@ def archive_transmitted_data_task(files): if instance_path.exists(): instance_path.replace(archive_dir / kind / instance_path.name) + + +def is_production(): + return bool(Variable.get("OKAPI_URL", "http://okapi-prod").find("prod") > 0) + + +def return_success_test_instance(files) -> dict: + return_files = [] + try: + return_files = [f for f in files["file_list"]] + return {"success": return_files, "failures": []} + except TypeError as e: + return_files = files + return {"success": files, "failures": []}