def execute(self, context: 'Context'):
hook = CloudAutoMLHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
result = hook.get_model(
model_id=self.model_id,
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return Model.to_dict(result)
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)
def execute(self, context: 'Context'):
hook = CloudAutoMLHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
self.log.info("Creating model.")
operation = hook.create_model(
model=self.model,
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
result = Model.to_dict(operation.result())
model_id = hook.extract_object_id(result)
self.log.info("Model created: %s", model_id)
self.xcom_push(context, key="model_id", value=model_id)
return result