diff --git a/test/data/test_time_window.py b/test/data/test_time_window.py index ad13a4d02..1e18405e0 100644 --- a/test/data/test_time_window.py +++ b/test/data/test_time_window.py @@ -85,6 +85,10 @@ def __call__(self, time, variable): }, ) + async def fetch(self, time, variable): + """Mock async data fetch.""" + return self(time, variable) + class TestTimeWindowInitialization: """Test TimeWindow initialization and validation.""" @@ -357,6 +361,34 @@ def test_time_coordinate_alignment(self): assert len(result.time) == 1 assert result.time.values[0] == np.datetime64(base_time) + @pytest.mark.asyncio + async def test_fetch_matches_sync_output_coordinates(self): + """Test async fetch applies offsets and preserves output coordinates.""" + ds = MockDataSource() + tw = TimeWindow( + datasource=ds, + offsets=[timedelta(hours=-6), timedelta(hours=0), timedelta(hours=6)], + suffixes=["_tm1", "_t", "_tp1"], + ) + + base_time = datetime(2024, 1, 1, 12, 0) + result = await tw.fetch(base_time, ["t2m", "u10m"]) + + expected_vars = [ + "t2m_tm1", + "t2m_t", + "t2m_tp1", + "u10m_tm1", + "u10m_t", + "u10m_tp1", + ] + actual_vars = [str(v) for v in result.coords["variable"].values] + assert actual_vars == expected_vars + assert result.time.values[0] == np.datetime64(base_time) + assert ds.call_history[0]["time"] == [datetime(2024, 1, 1, 6, 0)] # -6h + assert ds.call_history[1]["time"] == [datetime(2024, 1, 1, 12, 0)] # 0h + assert ds.call_history[2]["time"] == [datetime(2024, 1, 1, 18, 0)] # +6h + class TestTimeWindowErrorHandling: """Test error handling in TimeWindow."""