def test_transform_batch_with_existing_buffer(self, xr_mock, *args):
# Check that data in batch learning are concatenated
input_step = MagicMock()
input_step._should_stop.return_value = False
time = pd.date_range('2000-01-01', freq='1D', periods=7)
time2 = pd.date_range('2000-01-14', freq='1D', periods=7)
time3 = pd.date_range('2000-01-01', freq='1D', periods=14)
self.module_mock.transform.return_value = xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"],
coords={'time': time2})
step = Step(self.module_mock, {"x": input_step}, file_manager=MagicMock())
da = xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"], coords={'time': time})
step.buffer = {"test": da}
xr_mock.concat.return_value = xr.DataArray([2, 3, 4, 3, 3, 1, 2, 2, 3, 4, 3, 3, 1, 2], dims=["time"],
coords={'time': time3})
step.get_result(pd.Timestamp("2000.01.07"), pd.Timestamp("2020.01.14"))
# Two calls, once in should_stop and once in _transform
input_step.get_result.assert_has_calls(
[call(pd.Timestamp('2000-01-07 00:00:00'), pd.Timestamp('2020-01-14 00:00:00')),
call(pd.Timestamp('2000-01-07 00:00:00'), pd.Timestamp('2020-01-14 00:00:00'))])
xr_mock.concat.assert_called_once()
xr.testing.assert_equal(da, list(xr_mock.concat.call_args_list[0])[0][0][0])
xr.testing.assert_equal(xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"],
coords={'time': time2}), list(xr_mock.concat.call_args_list[0])[0][0][1])
assert {'dim': 'time'} == list(xr_mock.concat.call_args_list[0])[1]
def test_get_result(self):
# Tests if the get_result method calls correctly the previous step and the module
input_step = MagicMock()
input_step_result_mock = MagicMock()
input_step.get_result.return_value = input_step_result_mock
input_step._should_stop.return_value = False
time = pd.date_range('2000-01-01', freq='1H', periods=7)
self.module_mock.transform.return_value = xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"],
coords={'time': time})
step = Step(self.module_mock, {"x": input_step}, file_manager=MagicMock())
step.get_result(pd.Timestamp("2000.01.01"), pd.Timestamp("2020.12.12"))
# Two calls, once in should_stop and once in _transform
input_step.get_result.assert_has_calls(
[call(pd.Timestamp('2000-01-01 '), pd.Timestamp('2020-12-12 ')),
call(pd.Timestamp('2000-01-01 '), pd.Timestamp('2020-12-12 '))])
self.module_mock.transform.assert_called_once_with(x=input_step_result_mock)