Python MLEngineGetModelOperator.execute Examples

Example #1

File: test_mlengine.py Project: yyz940922/airflow_date_issue_corrected

    def test_success_get_model(self, mock_hook):
        task = MLEngineGetModelOperator(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)

        result = 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.get_model.assert_called_once_with(
            project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME)
        self.assertEqual(mock_hook.return_value.get_model.return_value, result)

Example #2

File: test_mlengine.py Project: ysktir/airflow-1

    def test_success_get_model(self, mock_hook):
        task = MLEngineGetModelOperator(
            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,
            impersonation_chain=TEST_IMPERSONATION_CHAIN,
        )

        result = task.execute(None)

        mock_hook.assert_called_once_with(
            delegate_to=TEST_DELEGATE_TO,
            gcp_conn_id=TEST_GCP_CONN_ID,
            impersonation_chain=TEST_IMPERSONATION_CHAIN,
        )
        mock_hook.return_value.get_model.assert_called_once_with(
            project_id=TEST_PROJECT_ID, model_name=TEST_MODEL_NAME)
        assert mock_hook.return_value.get_model.return_value == result