diff --git a/nowcast/next_workers.py b/nowcast/next_workers.py index ed874afe..7761da47 100644 --- a/nowcast/next_workers.py +++ b/nowcast/next_workers.py @@ -59,16 +59,15 @@ def after_download_weather(msg, config, checklist): "success 1km 12": [], } if msg.type.startswith("success"): + data_date = arrow.now().shift(days=-1).format("YYYY-MM-DD") if msg.type.endswith("2.5km 06"): - data_date = arrow.now().shift(days=-1).format("YYYY-MM-DD") - for data_src in config["rivers"]["stations"]: - for river_name in config["rivers"]["stations"][data_src]: - next_workers["success 2.5km 06"].append( - NextWorker( - "nowcast.workers.collect_river_data", - args=[data_src, river_name, "--data-date", data_date], - ) + for river_name in config["rivers"]["stations"]["ECCC"]: + next_workers["success 2.5km 06"].append( + NextWorker( + "nowcast.workers.collect_river_data", + args=["ECCC", river_name, "--data-date", data_date], ) + ) for stn in config["observations"]["ctd data"]["stations"]: next_workers["success 2.5km 06"].append( NextWorker("nowcast.workers.get_onc_ctd", args=[stn]) @@ -89,6 +88,13 @@ def after_download_weather(msg, config, checklist): race_condition_workers = {"grib_to_netcdf", "make_ssh_files"} return next_workers[msg.type], race_condition_workers if msg.type.endswith("2.5km 12"): + for river_name in config["rivers"]["stations"]["USGS"]: + next_workers["success 2.5km 12"].append( + NextWorker( + "nowcast.workers.collect_river_data", + args=["USGS", river_name, "--data-date", data_date], + ) + ) next_workers["success 2.5km 12"].extend( [ NextWorker("nowcast.workers.collect_NeahBay_ssh", args=["06"]), diff --git a/tests/test_next_workers.py b/tests/test_next_workers.py index 471ce99d..bbdc12f2 100644 --- a/tests/test_next_workers.py +++ b/tests/test_next_workers.py @@ -186,11 +186,6 @@ def mock_now(): monkeypatch.setattr(next_workers.arrow, "now", mock_now) - def mock_utcnow(): - return arrow.get("2018-12-27") - - monkeypatch.setattr(next_workers.arrow, "utcnow", mock_utcnow) - workers, race_condition_workers = next_workers.after_download_weather( Message("download_weather", "success 2.5km 06"), config, checklist ) @@ -210,11 +205,6 @@ def mock_utcnow(): ["ECCC", "Fraser", "--data-date", "2018-12-26"], host="localhost", ), - NextWorker( - "nowcast.workers.collect_river_data", - ["USGS", "SkagitMountVernon", "--data-date", "2018-12-26"], - host="localhost", - ), NextWorker("nowcast.workers.get_onc_ctd", ["SCVIP"], host="localhost"), NextWorker("nowcast.workers.get_onc_ctd", ["SEVIP"], host="localhost"), NextWorker("nowcast.workers.get_onc_ctd", ["USDDL"], host="localhost"), @@ -226,20 +216,28 @@ def mock_utcnow(): assert workers == expected assert race_condition_workers == {"grib_to_netcdf", "make_ssh_files"} - def test_success_2_5_km_12(self, config, checklist): + def test_success_2_5_km_12(self, config, checklist, monkeypatch): + def mock_now(): + return arrow.get("2018-12-27") + + monkeypatch.setattr(next_workers.arrow, "now", mock_now) + workers, race_condition_workers = next_workers.after_download_weather( Message("download_weather", "success 2.5km 12"), config, checklist ) expected = [ + NextWorker( + "nowcast.workers.collect_river_data", + ["USGS", "SkagitMountVernon", "--data-date", "2018-12-26"], + host="localhost", + ), NextWorker("nowcast.workers.collect_NeahBay_ssh", ["06"], host="localhost"), NextWorker( "nowcast.workers.grib_to_netcdf", ["nowcast+"], host="localhost" ), NextWorker("nowcast.workers.download_live_ocean", [], host="localhost"), ] - assert len(workers) == len(expected) - for next_worker in expected: - assert next_worker in workers + assert workers == expected assert race_condition_workers == { "grib_to_netcdf", "make_live_ocean_files", @@ -274,11 +272,6 @@ def mock_now(): monkeypatch.setattr(next_workers.arrow, "now", mock_now) - def mock_utcnow(): - return arrow.get("2018-12-27") - - monkeypatch.setattr(next_workers.arrow, "utcnow", mock_utcnow) - workers, race_condition_workers = next_workers.after_collect_weather( Message("collect_weather", "success 2.5km 06"), config, checklist ) @@ -298,11 +291,6 @@ def mock_utcnow(): ["ECCC", "Fraser", "--data-date", "2018-12-26"], host="localhost", ), - NextWorker( - "nowcast.workers.collect_river_data", - ["USGS", "SkagitMountVernon", "--data-date", "2018-12-26"], - host="localhost", - ), NextWorker("nowcast.workers.get_onc_ctd", ["SCVIP"], host="localhost"), NextWorker("nowcast.workers.get_onc_ctd", ["SEVIP"], host="localhost"), NextWorker("nowcast.workers.get_onc_ctd", ["USDDL"], host="localhost"), @@ -317,11 +305,21 @@ def mock_utcnow(): assert workers == expected assert race_condition_workers == {"grib_to_netcdf", "make_ssh_files"} - def test_success_2_5_km_12(self, config, checklist): + def test_success_2_5_km_12(self, config, checklist, monkeypatch): + def mock_now(): + return arrow.get("2018-12-27") + + monkeypatch.setattr(next_workers.arrow, "now", mock_now) + workers, race_condition_workers = next_workers.after_collect_weather( Message("collect_weather", "success 2.5km 12"), config, checklist ) expected = [ + NextWorker( + "nowcast.workers.collect_river_data", + ["USGS", "SkagitMountVernon", "--data-date", "2018-12-26"], + host="localhost", + ), NextWorker("nowcast.workers.collect_NeahBay_ssh", ["06"], host="localhost"), NextWorker( "nowcast.workers.grib_to_netcdf", ["nowcast+"], host="localhost" @@ -331,9 +329,7 @@ def test_success_2_5_km_12(self, config, checklist): "nowcast.workers.collect_weather", ["18", "2.5km"], host="localhost" ), ] - assert len(workers) == len(expected) - for next_worker in expected: - assert next_worker in workers + assert workers == expected assert race_condition_workers == { "grib_to_netcdf", "make_live_ocean_files",