def test_further_elements_already_buffered(self):
time = pd.date_range('2000-01-01', freq='24H', periods=7)
step = Step(self.module_mock, {"x": self.step_mock}, file_manager=MagicMock())
step.buffer = {"STEP": xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"], coords={'time': time})}
result = step.further_elements(pd.Timestamp("2000-01-05"))
self.step_mock.further_elements.assert_not_called()
self.assertEqual(result, True)
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_further_elements_target_false(self):
target_step = MagicMock()
target_step.further_elements.return_value = False
time = pd.date_range('2000-01-01', freq='1H', periods=7)
step = Step(self.module_mock, {"x": self.step_mock}, targets={"target": target_step}, file_manager=MagicMock())
step.buffer = {"STEP": xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"], coords={'time': time})}
result = step.further_elements(pd.Timestamp("2000.12.12"))
target_step.further_elements.assert_called_once_with(pd.Timestamp("2000.12.12"))
self.assertFalse(result)
def test_reset(self):
step = Step(MagicMock(), MagicMock(), MagicMock())
step.buffer = MagicMock()
step.current_run_setting = RunSetting(computation_mode=ComputationMode.Transform)
step.finished = True
step.reset()
self.assertIsNone(None)
assert step.current_run_setting.computation_mode == ComputationMode.Default
assert step._should_stop(None, None) == False
assert step.finished == False
def test_reset(self):
step = Step(MagicMock(), MagicMock(), MagicMock())
step.buffer = MagicMock()
step.computation_mode = ComputationMode.Transform
step.stop = True
step.finished = True
step.reset()
self.assertIsNone(None)
assert step.computation_mode == ComputationMode.Default
assert step.stop == False
assert step.finished == False