Skip to content

Commit

Permalink
Updates ftp transmission task to use the gather_files dict.
Browse files Browse the repository at this point in the history
  • Loading branch information
shelleydoljack committed May 15, 2024
1 parent f02ac4b commit ad613d5
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 2 deletions.
4 changes: 2 additions & 2 deletions libsys_airflow/plugins/data_exports/transmission_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ def transmit_data_http_task(gather_files, **kwargs) -> dict:


@task(multiple_outputs=True)
def transmit_data_ftp_task(conn_id, local_files) -> dict:
def transmit_data_ftp_task(conn_id, gather_files) -> dict:
"""
Transmit the data via ftp
Returns lists of files successfully transmitted and failures
Expand All @@ -87,7 +87,7 @@ def transmit_data_ftp_task(conn_id, local_files) -> dict:
remote_path = connection.extra_dejson["remote_path"]
success = []
failures = []
for f in local_files:
for f in gather_files["file_list"]:
remote_file_path = f"{remote_path}/{Path(f).name}"
try:
logger.info(f"Start transmission of file {f}")
Expand Down
28 changes: 28 additions & 0 deletions tests/data_exports/test_transmission_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from libsys_airflow.plugins.data_exports.transmission_tasks import (
gather_files_task,
transmit_data_http_task,
transmit_data_ftp_task,
archive_transmitted_data_task,
)

Expand Down Expand Up @@ -60,6 +61,19 @@ def mock_httpx_connection():
)


@pytest.fixture
def mock_ftphook_connection():
return Connection(
conn_id="ftp-example.com",
conn_type="ftp",
host="ftp://www.example.com",
login="username",
password="pass",
extra={"remote_path": "/remote/path/dir"},
schema="ftp",
)


@pytest.fixture
def mock_httpx_success():
return httpx.Client(
Expand Down Expand Up @@ -88,6 +102,20 @@ def test_gather_full_dump_files(mocker):
assert marc_files["s3"]


def test_transmit_data_ftp_task(mocker, mock_ftphook_connection, mock_marc_files, caplog):
ftp_hook = mocker.patch("airflow.providers.ftp.hooks.ftp.FTPHook.store_file", return_value=True)
mocker.patch(
"libsys_airflow.plugins.data_exports.transmission_tasks.Connection.get_connection_from_secrets",
return_value=mock_ftphook_connection,
)

transmit_data = transmit_data_ftp_task.function("ftp-example.com", mock_marc_files)
assert len(transmit_data["success"]) == 3
assert "Start transmission of file" in caplog.text
assert ftp_hook.store_file.called_with(
"/remote/path/dir/2024022914.mrc", "2024022914.mrc"
)

def test_transmit_data_task(
mocker, mock_httpx_connection, mock_httpx_success, mock_marc_files, caplog
):
Expand Down

0 comments on commit ad613d5

Please sign in to comment.