def test_fit_with_targets(self): input_dict = {'input_data': None} target_mock = MagicMock() step = Step(self.module_mock, self.step_mock, MagicMock(), targets={"target": MagicMock()}) step._fit(input_dict, {"target": target_mock}) self.module_mock.fit.assert_called_once_with(**input_dict, target=target_mock)
def test_fit(self): input_dict = {'input_data': None} step = Step(self.module_mock, self.step_mock, file_manager=MagicMock()) step._fit(input_dict, {}) self.module_mock.fit.assert_called_once_with(**input_dict)
def test_transform(self): input_dict = {'input_data': None} step = Step(self.module_mock, {"x": self.step_mock}, None) step._fit(input_dict, {}) step._transform(input_dict) self.module_mock.transform.assert_called_once_with(**input_dict)