def test_execute(self, mock_hook):
mock_hook.return_value.get_model.return_value = Model(name=MODEL_PATH)
mock_hook.return_value.extract_object_id = extract_object_id
op = AutoMLGetModelOperator(
model_id=MODEL_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook.return_value.get_model.assert_called_once_with(
location=GCP_LOCATION,
metadata=None,
model_id=MODEL_ID,
project_id=GCP_PROJECT_ID,
retry=None,
timeout=None,
)
def test_execute(self, mock_hook, mock_xcom):
mock_hook.return_value.create_model.return_value.result.return_value = Model(
name=MODEL_PATH)
mock_hook.return_value.extract_object_id = extract_object_id
op = AutoMLTrainModelOperator(
model=MODEL,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook.return_value.create_model.assert_called_once_with(
model=MODEL,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
retry=None,
timeout=None,
metadata=None,
)
mock_xcom.assert_called_once_with(None, key="model_id", value=MODEL_ID)