def test_success(self, mock_hook): task = MLEngineListVersionsOperator( task_id="task-id", project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME, gcp_conn_id=TEST_GCP_CONN_ID, delegate_to=TEST_DELEGATE_TO, ) task.execute(None) mock_hook.assert_called_once_with(delegate_to=TEST_DELEGATE_TO, gcp_conn_id=TEST_GCP_CONN_ID) mock_hook.return_value.list_versions.assert_called_once_with( project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME, )
def test_missing_model_name(self): with self.assertRaises(AirflowException): MLEngineListVersionsOperator( task_id="task-id", project_id=TEST_PROJECT_ID, model_name=None, gcp_conn_id=TEST_GCP_CONN_ID, delegate_to=TEST_DELEGATE_TO, )
"runtime_version": "1.14", "machineType": "mls1-c1-m2", "framework": "TENSORFLOW", "pythonVersion": "3.5" }) set_defaults_version = MLEngineSetDefaultVersionOperator( task_id="set-default-version", project_id=PROJECT_ID, model_name=MODEL_NAME, version_name="v2", ) list_version = MLEngineListVersionsOperator( task_id="list-version", project_id=PROJECT_ID, model_name=MODEL_NAME, ) list_version_result = BashOperator( bash_command="echo \"{{ task_instance.xcom_pull('list-version') }}\"", task_id="list-version-result", ) prediction = MLEngineBatchPredictionOperator( task_id="prediction", project_id="polidea-airflow", job_id="prediciton-{{ ts_nodash }}-{{ params.model_name }}", region="us-central1", model_name=MODEL_NAME, data_format="TEXT",