diff --git a/pyproject.toml b/pyproject.toml index 86c587c5..7aba605e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,7 @@ dependencies = [ "openpyxl", "pandas", "panel", - "prefect[shell,dask]", + "prefect[shell,dask]==3.8.4", "PyYAML", "s3fs", "shapely", diff --git a/src/echodataflow/deployment/flow_registry.py b/src/echodataflow/deployment/flow_registry.py index 92c9822d..5425e920 100644 --- a/src/echodataflow/deployment/flow_registry.py +++ b/src/echodataflow/deployment/flow_registry.py @@ -18,6 +18,10 @@ class FlowRegistration: entrypoint="echodataflow/flows/flows_acoustics.py:flow_raw2Sv", description="Incrementally convert newly available raw sonar files to Sv.", ), + "raw2Sv_CPS": FlowRegistration( + entrypoint="echodataflow/flows/flows_acoustics.py:flow_raw2Sv_CPS", + description="Convert newly registered RAW files to Sv using the CPS processing database.", + ), "create_MVBS": FlowRegistration( entrypoint="echodataflow/flows/flows_acoustics.py:flow_create_MVBS", ), @@ -63,9 +67,6 @@ class FlowRegistration: "copy_trawl": FlowRegistration( entrypoint="echodataflow/flows/flows_simulation.py:flow_copy_trawl", ), - "simulate_transects": FlowRegistration( - entrypoint="echodataflow/flows/flows_simulation.py:flow_simulate_transects", - ), "update_cache_MVBS": FlowRegistration( entrypoint="echodataflow/flows/flows_viz_cloud.py:flow_update_cache_MVBS", ), diff --git a/src/echodataflow/flows/flows_CPS.py b/src/echodataflow/flows/flows_CPS.py index 4c8b9846..b80c313e 100644 --- a/src/echodataflow/flows/flows_CPS.py +++ b/src/echodataflow/flows/flows_CPS.py @@ -1,4 +1,3 @@ -import asyncio from pathlib import Path import dask_image.ndfilters @@ -7,9 +6,7 @@ import numpy as np import pandas as pd import xarray as xr -from prefect import flow, get_client, runtime -from prefect.states import Cancelled -from echodataflow.flows.flows_helper import deployment_already_running +from prefect import flow from prefect_dask import DaskTaskRunner from echodataflow.utils.processing_ledger import get_completed_sv_files, resolve_database @@ -255,30 +252,9 @@ def flow_process_CPS( range_bin: str = "10m", dist_bin: str = "0.5nmi", nasc_process_id: int = 1928, + exclude_before: str | None = None, ): - # Prevent overlapping runs of this deployment - already_running = asyncio.run( - deployment_already_running() - ) - - if already_running: - - async def cancel_run(): - async with get_client() as client: - await client.set_flow_run_state( - flow_run_id=runtime.flow_run.id, - state=Cancelled( - message=( - "Another instance of this " - "flow is already running" - ) - ), - ) - - asyncio.run(cancel_run()) - return - path_main = Path(path_main) path_transect = Path( @@ -323,7 +299,7 @@ async def cancel_run(): ) # --------------------------------------------- - # Find completed transects still needing CPS + # Select eligible completed transects # --------------------------------------------- current = pd.read_csv( @@ -336,8 +312,38 @@ async def cancel_run(): }, ) + eligible = current.copy() + + eligible["transectStart"] = pd.to_datetime( + eligible["transectStart"], + utc=True, + errors="coerce", + ) + + eligible["transectEnd"] = pd.to_datetime( + eligible["transectEnd"], + utc=True, + errors="coerce", + ) + + # Ignore historical transects that started + # before this deployment's processing window. + if exclude_before is not None: + cutoff = pd.Timestamp( + exclude_before + ) + + if cutoff.tzinfo is None: + cutoff = cutoff.tz_localize("UTC") + else: + cutoff = cutoff.tz_convert("UTC") + + eligible = eligible.loc[ + eligible["transectStart"] >= cutoff + ].copy() + # Ignore transects that have not finished yet. - completed = current.dropna( + completed = eligible.dropna( subset=[ "transectPart", "transectStart", diff --git a/src/echodataflow/flows/flows_acoustics.py b/src/echodataflow/flows/flows_acoustics.py index 5d8c0152..7e214aa9 100644 --- a/src/echodataflow/flows/flows_acoustics.py +++ b/src/echodataflow/flows/flows_acoustics.py @@ -37,6 +37,15 @@ plan_mvbs_slices, read_or_create_ledger, ) + +from echodataflow.utils.processing_ledger import ( + get_raw_files_to_process, + initialize_ledger, + mark_raw_completed, + mark_raw_failed, + mark_raw_processing, + resolve_database, +) from echodataflow.operations.operations_storage import S3CopySettings, S3CopyWorkItem from echodataflow.tasks.tasks_acoustics import ( task_create_MVBS, @@ -142,7 +151,7 @@ def flow_raw2Sv( f"Limiting to first {new_file_num_limit} files." ) new_files = new_files[:new_file_num_limit] - print(f"Files to process: \n" + "".join([f"- {nf}\n" for nf in new_files])) + print("Files to process: \n" + "".join([f"- {nf}\n" for nf in new_files])) settings = RawToSvSettings( output_directory=path_Sv_zarr, @@ -227,6 +236,179 @@ async def set_failed_state(): asyncio.run(set_failed_state()) raise Exception(error_msg) + +@flow(log_prints=True, task_runner=dask_task_runner_from_environment()) +def flow_raw2Sv_CPS( + exclude_before: str | None = None, + exclude_raw_file: list[str] = [], + parallel: bool = False, + encode_mode: str = "power", + waveform_mode: str = "CW", + depth_offset: float = 9.5, + sonar_model: str = "EK80", + datagram_type: str | None = None, + nmea_sentence: str | None = None, + path_main: str = "", + processing_db: str = "processing.db", + new_file_num_limit: int = 50, + add_depth: bool = True, + add_location: bool = True, + add_splitbeam_angle: bool = False, +): + """Convert newly registered RAW files to Sv using the CPS processing database.""" + + path_Sv_zarr = Path(path_main) / "Sv" + path_Sv_zarr.mkdir(parents=True, exist_ok=True) + + db_path = resolve_database(path_main, processing_db) + initialize_ledger(db_path) + + raw_files = get_raw_files_to_process( + db_path, + limit=new_file_num_limit, + ) + + if exclude_before is not None: + exclude_before_dt = pd.to_datetime( + exclude_before, + utc=True, + ).to_pydatetime() + + raw_files = [ + raw_path + for raw_path in raw_files + if extract_datetime_from_filename(raw_path.name) + >= exclude_before_dt + ] + + if exclude_raw_file: + excluded = set(exclude_raw_file) + raw_files = [ + raw_path + for raw_path in raw_files + if raw_path.name not in excluded + ] + + print(f"Found {len(raw_files)} RAW files to process") + print( + "Files to process:\n" + + "".join(f"- {raw_path.name}\n" for raw_path in raw_files) + ) + + if not raw_files: + return + + settings = RawToSvSettings( + output_directory=str(path_Sv_zarr), + encode_mode=encode_mode, + waveform_mode=waveform_mode, + depth_offset=depth_offset, + sonar_model=sonar_model, + datagram_type=datagram_type, + nmea_sentence=nmea_sentence, + add_depth=add_depth, + add_location=add_location, + add_splitbeam_angle=add_splitbeam_angle, + ) + + errors = [] + + if parallel: + print("Processing RAW files in parallel") + futures = {} + + for raw_path in raw_files: + mark_raw_processing(db_path, raw_path) + + future = task_raw2Sv.with_options( + task_run_name=raw_path.name, + name=raw_path.name, + retries=3, + ).submit( + RawToSvWorkItem(raw_path=str(raw_path)), + settings, + ) + + futures[future] = raw_path + + for future in as_completed(futures): + raw_path = futures[future] + + try: + result = future.result() + + mark_raw_completed( + db_path, + raw_path, + result.filename_Sv, + result.first_ping_time, + result.last_ping_time, + ) + + except Exception as exc: + mark_raw_failed( + db_path, + raw_path, + str(exc), + ) + errors.append(exc) + print(f"Error converting {raw_path.name}: {exc}") + + else: + print("Processing RAW files sequentially") + + for raw_path in raw_files: + try: + print(f"Converting {raw_path.name}") + + mark_raw_processing( + db_path, + raw_path, + ) + + result = task_raw2Sv.with_options( + task_run_name=raw_path.name, + name=raw_path.name, + retries=3, + )( + RawToSvWorkItem(raw_path=str(raw_path)), + settings, + ) + + mark_raw_completed( + db_path, + raw_path, + result.filename_Sv, + result.first_ping_time, + result.last_ping_time, + ) + + except Exception as exc: + mark_raw_failed( + db_path, + raw_path, + str(exc), + ) + errors.append(exc) + print(f"Error converting {raw_path.name}: {exc}") + + if errors: + error_msg = ( + f"{len(errors)} errors during raw to Sv conversion " + f"out of {len(raw_files)} files" + ) + + async def set_failed_state(): + async with get_client() as client: + await client.set_flow_run_state( + flow_run_id=runtime.flow_run.id, + state=Failed(message=error_msg), + ) + + asyncio.run(set_failed_state()) + raise RuntimeError(error_msg) + + @flow(log_prints=True) async def flow_create_MVBS( time_offset_seconds: float = 0.0, diff --git a/src/echodataflow/flows/flows_simulation.py b/src/echodataflow/flows/flows_simulation.py index 1e3b8752..081d98b4 100644 --- a/src/echodataflow/flows/flows_simulation.py +++ b/src/echodataflow/flows/flows_simulation.py @@ -246,11 +246,12 @@ def flow_simulate_transects( transect_state_key = _var_key(prefix="transect_state") state = Variable.get(transect_state_key, default=None) - transect_num_curr = ( - start_transect_num - if state is None - else int(state) - ) + if state is None: + transect_num_curr = start_transect_num + action = "open" + else: + transect_num_str, action = str(state).split(":", maxsplit=1) + transect_num_curr = int(transect_num_str) if transect_num_curr > max_transects: print("All simulated transects have been generated.") diff --git a/src/echodataflow/flows/flows_viz_cloud.py b/src/echodataflow/flows/flows_viz_cloud.py index 7e472b2f..6158381e 100644 --- a/src/echodataflow/flows/flows_viz_cloud.py +++ b/src/echodataflow/flows/flows_viz_cloud.py @@ -9,7 +9,7 @@ from prefect import flow, get_run_logger -from echodataflow.utils.utils import round_up_mins, get_slice_start_end_times +from echodataflow.utils.utils import get_slice_start_end_times @flow() @@ -256,10 +256,10 @@ def flow_update_cache_CPS( # Find latest completed CPS transect # ----------------------------------------------------- - cps_files = sorted( - path_CPS.glob("transect_*_CPS.zarr"), - key=lambda path: path.stat().st_mtime, - ) + def _transect_number(path: Path) -> int: + return int(path.name.replace("transect_", "").replace("_CPS.zarr", "")) + + cps_files = sorted(path_CPS.glob("transect_*_CPS.zarr"), key=_transect_number) if not cps_files: print( diff --git a/src/echodataflow/operations/operations_watchdog.py b/src/echodataflow/operations/operations_watchdog.py index e49c0cbc..8649d5f9 100644 --- a/src/echodataflow/operations/operations_watchdog.py +++ b/src/echodataflow/operations/operations_watchdog.py @@ -1,6 +1,8 @@ from pathlib import Path from prefect.events import emit_event +from prefect.events.worker import EventsWorker +from watchdog.observers import Observer from echodataflow.utils.file_watcher import watch_directory, watch_file from echodataflow.utils.processing_ledger import ( @@ -9,55 +11,44 @@ ) -RAW_UPDATE_EVENT = "echodataflow.raw.updated" -RAW_RESOURCE_ID = "raw-monitor" - TRANSECT_UPDATE_EVENT = "echodataflow.transect.updated" TRANSECT_RESOURCE_ID = "transect-start-end-time" TRANSECT_RELATED_RESOURCE_ID = "transect-monitor" -def emit_raw_update_event(path: Path) -> None: - """Emit a Prefect event when a RAW file arrives.""" - emit_event( - event=RAW_UPDATE_EVENT, - resource={ - "prefect.resource.id": RAW_RESOURCE_ID, - "prefect.resource.name": RAW_RESOURCE_ID, - "path": str(path), - }, - ) +def _flush_events() -> None: + """Wait until queued Prefect events have been sent.""" + EventsWorker.instance().wait_until_empty() -def register_and_emit_raw_update(path: Path, db_path: str | Path) -> None: - """Register a RAW file in the ledger, then emit its Prefect event.""" +def register_raw_update( + path: Path, + db_path: str | Path, +) -> None: + """Register a RAW file change in the processing ledger.""" register_raw_file(db_path, path) - emit_raw_update_event(path) def watch_raw_directory( path: str | Path, db_path: str | Path, -): - """Watch a directory for new RAW files.""" +) -> Observer: + """Watch a directory for RAW files and keep the processing ledger updated.""" raw_directory = Path(path).resolve() initialize_ledger(db_path) - # Reconcile files that already exist when the watcher starts - existing_raw_files = list(raw_directory.glob("*.raw")) - - for raw_path in existing_raw_files: + # Reconcile files that already exist when the watcher starts. + for raw_path in raw_directory.glob("*.raw"): register_raw_file(db_path, raw_path) - # Wake raw2Sv once after reconciliation. - if existing_raw_files: - emit_raw_update_event(raw_directory) - return watch_directory( directory=raw_directory, - callback=lambda raw_path: register_and_emit_raw_update(raw_path, db_path), + callback=lambda raw_path: register_raw_update( + raw_path, + db_path, + ), pattern="*.raw", ) @@ -65,7 +56,7 @@ def watch_raw_directory( def emit_transect_update_event(path: Path) -> None: """Emit a Prefect event when the transect CSV is updated.""" - emit_event( + event = emit_event( event=TRANSECT_UPDATE_EVENT, resource={ "prefect.resource.id": TRANSECT_RESOURCE_ID, @@ -80,6 +71,9 @@ def emit_transect_update_event(path: Path) -> None: ], ) + if event is not None: + _flush_events() + def watch_transect_file(path: str | Path): """Watch the transect start/end CSV and emit a Prefect event on update.""" @@ -87,4 +81,4 @@ def watch_transect_file(path: str | Path): return watch_file( target_file=path, callback=emit_transect_update_event, - ) + ) \ No newline at end of file diff --git a/src/echodataflow/services/viz_echogram_track_cps.py b/src/echodataflow/services/viz_echogram_track_cps.py index 8b4806d6..826d52f9 100644 --- a/src/echodataflow/services/viz_echogram_track_cps.py +++ b/src/echodataflow/services/viz_echogram_track_cps.py @@ -299,7 +299,7 @@ def plot_nasc( ds_nasc: xr.Dataset, title: str, ): - """Plot NASC as vertical bars along distance.""" + """Plot depth-integrated NASC along transect time.""" if "NASC" not in ds_nasc: return pn.pane.Markdown( @@ -310,11 +310,9 @@ def plot_nasc( # Select the channel closest to the configured target frequency. if "channel" in nasc.dims: - target_channel = ( - pick_channel_by_frequency( - ds_nasc, - TARGET_FREQUENCY, - ) + target_channel = pick_channel_by_frequency( + ds_nasc, + TARGET_FREQUENCY, ) nasc = nasc.sel( @@ -326,50 +324,62 @@ def plot_nasc( frequency_nominal=0 ) - # Remove singleton dimensions nasc = nasc.squeeze( drop=True ) print( - "NASC dims:", + "NASC dims before integration:", nasc.dims, ) print( - "NASC shape:", + "NASC shape before integration:", nasc.shape, ) - print( - "NASC coords:", - list( - nasc.coords - ), - ) - - # NASC should ultimately be one value per horizontal interval. - # If another dimension remains, integrate/sum over it for plotting. - while nasc.ndim > 1: - dim_to_reduce = ( - nasc.dims[0] - ) + # NASC is depth-resolved in the stored product: + # distance x depth. + # + # For the dashboard, integrate explicitly over depth so that + # one NASC value remains for each horizontal interval. + if "depth" in nasc.dims: nasc = nasc.sum( - dim=dim_to_reduce, + dim="depth", skipna=True, ) - dim = nasc.dims[0] + nasc = nasc.squeeze( + drop=True + ) + + print( + "NASC dims after integration:", + nasc.dims, + ) + print( + "NASC shape after integration:", + nasc.shape, + ) + + # Use the representative ping time associated with each + # horizontal NASC interval so the plot aligns with the echogram. + if "ping_time" in ds_nasc: + x = pd.to_datetime( + ds_nasc["ping_time"].values + ) + xlabel = "Time" + kdim = "ping_time" - if "distance" in nasc.coords: - x = nasc[ - "distance" - ].values + elif "distance" in nasc.coords: + x = nasc["distance"].values xlabel = "Distance (nmi)" + kdim = "distance" + else: - x = nasc[ - dim - ].values + dim = nasc.dims[0] + x = nasc[dim].values xlabel = dim + kdim = dim curve = hv.Curve( ( @@ -377,7 +387,7 @@ def plot_nasc( nasc.values, ), kdims=[ - xlabel, + kdim, ], vdims=[ "NASC", @@ -400,7 +410,6 @@ def plot_nasc( ylabel="NASC", ) - # --------------------------------------------------------------------- # Latest transect plotting # --------------------------------------------------------------------- diff --git a/src/echodataflow/utils/processing_ledger.py b/src/echodataflow/utils/processing_ledger.py index 01be1134..6c2aa0cb 100644 --- a/src/echodataflow/utils/processing_ledger.py +++ b/src/echodataflow/utils/processing_ledger.py @@ -15,6 +15,8 @@ select, update, ) +from sqlalchemy.dialects.postgresql import insert as postgresql_insert +from sqlalchemy.dialects.sqlite import insert as sqlite_insert from sqlalchemy.engine import Engine from sqlalchemy.sql import func @@ -115,8 +117,12 @@ def initialize_ledger(db_path: str | Path) -> None: def register_raw_file( db_path: str | Path, raw_path: str | Path, -) -> None: - """Register a RAW file, re-queueing it only when its contents changed.""" +) -> bool: + """Register a RAW file, re-queueing it only when its contents changed. + + Returns True when the ledger changed and downstream processing should + be notified, otherwise False. + """ raw_path = Path(raw_path) stat = raw_path.stat() @@ -132,16 +138,46 @@ def register_raw_file( ).first() if existing is None: - conn.execute( - insert(raw_sv).values( - raw_path=str(raw_path), - raw_filename=raw_path.name, - file_size=stat.st_size, - file_mtime_ns=stat.st_mtime_ns, - status="pending", + values = { + "raw_path": str(raw_path), + "raw_filename": raw_path.name, + "file_size": stat.st_size, + "file_mtime_ns": stat.st_mtime_ns, + "status": "pending", + } + + if engine.dialect.name == "postgresql": + stmt = ( + postgresql_insert(raw_sv) + .values(**values) + .on_conflict_do_nothing( + index_elements=[raw_sv.c.raw_path] + ) ) - ) - return + elif engine.dialect.name == "sqlite": + stmt = ( + sqlite_insert(raw_sv) + .values(**values) + .on_conflict_do_nothing( + index_elements=[raw_sv.c.raw_path] + ) + ) + else: + stmt = insert(raw_sv).values(**values) + + result = conn.execute(stmt) + + if result.rowcount == 1: + return True + + # Another callback registered the same RAW between our + # SELECT and INSERT. Re-read it instead of failing. + existing = conn.execute( + select( + raw_sv.c.file_size, + raw_sv.c.file_mtime_ns, + ).where(raw_sv.c.raw_path == str(raw_path)) + ).one() if ( existing.file_size != stat.st_size @@ -161,7 +197,9 @@ def register_raw_file( updated_at=func.current_timestamp(), ) ) + return True + return False def get_raw_files_to_process( db_path: str | Path, diff --git a/tests/test_flow_raw2Sv.py b/tests/test_flow_raw2Sv.py index 3298596d..2e3fb92a 100644 --- a/tests/test_flow_raw2Sv.py +++ b/tests/test_flow_raw2Sv.py @@ -1,10 +1,6 @@ import pytest from types import SimpleNamespace -pytestmark = pytest.mark.skip( - reason="Temporarily disabled while flows_acoustics.py is reverted to 3c1b9d3" -) - from echodataflow.flows import flows_acoustics @@ -65,7 +61,7 @@ def __call__(self, work_item, settings): FakeTask(), ) - flows_acoustics.flow_raw2Sv.fn( + flows_acoustics.flow_raw2Sv_CPS.fn( path_main=str(tmp_path), new_file_num_limit=1, ) @@ -148,7 +144,7 @@ async def set_flow_run_state(self, **kwargs): ) with pytest.raises(RuntimeError, match="1 errors during raw to Sv conversion"): - flows_acoustics.flow_raw2Sv.fn( + flows_acoustics.flow_raw2Sv_CPS.fn( path_main=str(tmp_path), new_file_num_limit=1, ) diff --git a/tests/test_flow_simulate_transects.py b/tests/test_flow_simulate_transects.py index 09f80dfc..09a0d5ab 100644 --- a/tests/test_flow_simulate_transects.py +++ b/tests/test_flow_simulate_transects.py @@ -34,23 +34,23 @@ def test_flow_simulate_transects_opens_closes_and_advances(monkeypatch, tmp_path "max_transects": 2, } - # First run: write complete transect 001. + # First run: open transect 001. flows_simulation.flow_simulate_transects.fn(**kwargs) df = pd.read_csv(transect_csv, dtype="string") assert df["transectPart"].tolist() == ["001"] assert df.loc[0, "transectStart"] == "2024-07-07T00:00:00+00:00" - assert df.loc[0, "transectEnd"] == "2024-07-07T00:10:00+00:00" + assert pd.isna(df.loc[0, "transectEnd"]) - # Second run: append complete transect 002. + # Second run: close transect 001. flows_simulation.flow_simulate_transects.fn(**kwargs) df = pd.read_csv(transect_csv, dtype="string") - assert df["transectPart"].tolist() == ["001", "002"] - assert df.loc[1, "transectStart"] == "2024-07-07T00:10:00+00:00" - assert df.loc[1, "transectEnd"] == "2024-07-07T00:20:00+00:00" + assert df["transectPart"].tolist() == ["001"] + assert df.loc[0, "transectStart"] == "2024-07-07T00:00:00+00:00" + assert df.loc[0, "transectEnd"] == "2024-07-07T00:10:00+00:00" def test_flow_simulate_transects_stops_after_maximum(monkeypatch, tmp_path, capsys): diff --git a/tests/test_operations_watchdog.py b/tests/test_operations_watchdog.py index 1c9b4f72..e7883971 100644 --- a/tests/test_operations_watchdog.py +++ b/tests/test_operations_watchdog.py @@ -1,29 +1,6 @@ from echodataflow.operations import operations_watchdog -def test_emit_raw_update_event(monkeypatch, tmp_path): - raw_file = tmp_path / "example.raw" - raw_file.touch() - - emitted = {} - - def fake_emit_event(**kwargs): - emitted.update(kwargs) - - monkeypatch.setattr( - operations_watchdog, - "emit_event", - fake_emit_event, - ) - - operations_watchdog.emit_raw_update_event(raw_file) - - assert emitted["event"] == "echodataflow.raw.updated" - assert emitted["resource"]["prefect.resource.id"] == "raw-monitor" - assert emitted["resource"]["prefect.resource.name"] == "raw-monitor" - assert emitted["resource"]["path"] == str(raw_file) - - def test_watch_raw_directory(monkeypatch, tmp_path): called = {} db_path = tmp_path / "processing.db" @@ -59,12 +36,12 @@ def fake_watch_directory( ) assert result == "observer" - assert called["directory"] == tmp_path + assert called["directory"] == tmp_path.resolve() assert called["db_path"] == db_path assert called["pattern"] == "*.raw" -def test_register_and_emit_raw_update(monkeypatch, tmp_path): +def test_register_raw_update(monkeypatch, tmp_path): raw_file = tmp_path / "example.raw" db_path = tmp_path / "processing.db" @@ -73,40 +50,33 @@ def test_register_and_emit_raw_update(monkeypatch, tmp_path): def fake_register_raw_file(db, path): called["registered"] = (db, path) - def fake_emit_raw_update_event(path): - called["emitted"] = path - monkeypatch.setattr( operations_watchdog, "register_raw_file", fake_register_raw_file, ) - monkeypatch.setattr( - operations_watchdog, - "emit_raw_update_event", - fake_emit_raw_update_event, - ) - - operations_watchdog.register_and_emit_raw_update( + operations_watchdog.register_raw_update( raw_file, db_path, ) assert called["registered"] == (db_path, raw_file) - assert called["emitted"] == raw_file -def test_watch_raw_directory_reconciles_existing_raw_files(monkeypatch, tmp_path): +def test_watch_raw_directory_reconciles_existing_raw_files( + monkeypatch, + tmp_path, +): raw_a = tmp_path / "a.raw" raw_b = tmp_path / "b.raw" + raw_a.touch() raw_b.touch() db_path = tmp_path / "processing.db" registered = [] - emitted = [] monkeypatch.setattr( operations_watchdog, @@ -114,12 +84,6 @@ def test_watch_raw_directory_reconciles_existing_raw_files(monkeypatch, tmp_path lambda db, path: registered.append((db, path)), ) - monkeypatch.setattr( - operations_watchdog, - "emit_raw_update_event", - emitted.append, - ) - monkeypatch.setattr( operations_watchdog, "watch_directory", @@ -132,11 +96,11 @@ def test_watch_raw_directory_reconciles_existing_raw_files(monkeypatch, tmp_path ) assert result == "observer" + assert {path for _, path in registered} == { raw_a.resolve(), raw_b.resolve(), } - assert emitted == [tmp_path.resolve()] def test_emit_transect_update_event(monkeypatch, tmp_path): @@ -157,5 +121,8 @@ def fake_emit_event(**kwargs): operations_watchdog.emit_transect_update_event(target) assert emitted["event"] == "echodataflow.transect.updated" - assert emitted["resource"]["prefect.resource.id"] == "transect-start-end-time" - assert emitted["resource"]["path"] == str(target) + assert ( + emitted["resource"]["prefect.resource.id"] + == "transect-start-end-time" + ) + assert emitted["resource"]["path"] == str(target) \ No newline at end of file diff --git a/tests/test_processing_ledger.py b/tests/test_processing_ledger.py index 28f310e6..ec2bcea7 100644 --- a/tests/test_processing_ledger.py +++ b/tests/test_processing_ledger.py @@ -11,6 +11,7 @@ register_raw_file, resolve_database, ) +from concurrent.futures import ThreadPoolExecutor def test_database_url_from_path(tmp_path): @@ -73,8 +74,9 @@ def test_register_raw_file(tmp_path): raw_path.touch() initialize_ledger(db_path) - register_raw_file(db_path, raw_path) - register_raw_file(db_path, raw_path) + + assert register_raw_file(db_path, raw_path) is True + assert register_raw_file(db_path, raw_path) is False with sqlite3.connect(db_path) as conn: rows = conn.execute( @@ -234,7 +236,8 @@ def test_register_raw_file_requeues_changed_file(tmp_path): raw_path.write_bytes(b"first") initialize_ledger(db_path) - register_raw_file(db_path, raw_path) + + assert register_raw_file(db_path, raw_path) is True mark_raw_completed( db_path, @@ -245,7 +248,8 @@ def test_register_raw_file_requeues_changed_file(tmp_path): ) raw_path.write_bytes(b"first plus more data") - register_raw_file(db_path, raw_path) + + assert register_raw_file(db_path, raw_path) is True with sqlite3.connect(db_path) as conn: row = conn.execute( @@ -268,4 +272,33 @@ def test_register_raw_file_large_mtime(tmp_path): register_raw_file(db_path, raw_path) # A nanosecond mtime is much larger than a 32-bit integer. - assert raw_path.stat().st_mtime_ns > 2**31 \ No newline at end of file + assert raw_path.stat().st_mtime_ns > 2**31 + +def test_register_raw_file_concurrently_is_idempotent(tmp_path): + db_path = tmp_path / "processing.db" + raw_path = tmp_path / "test.raw" + + raw_path.write_bytes(b"raw-data") + initialize_ledger(db_path) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(register_raw_file, db_path, raw_path) + for _ in range(2) + ] + + results = [future.result() for future in futures] + + assert sorted(results) == [False, True] + + with sqlite3.connect(db_path) as conn: + rows = conn.execute( + """ + SELECT raw_path, raw_filename, status + FROM raw_sv + """ + ).fetchall() + + assert rows == [ + (str(raw_path), "test.raw", "pending"), + ] \ No newline at end of file