def testSuccessCreateVersion(self):
        with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
                as mock_hook:
            success_response = {'name': 'some-name', 'done': True}
            hook_instance = mock_hook.return_value
            hook_instance.create_version.return_value = success_response

            training_op = MLEngineVersionOperator(version=self.VERSION_INPUT,
                                                  **self.VERSION_DEFAULT_ARGS)
            training_op.execute(None)

            mock_hook.assert_called_with(gcp_conn_id='google_cloud_default',
                                         delegate_to=None)
            # Make sure only 'create_version' is invoked on hook instance
            self.assertEquals(len(hook_instance.mock_calls), 1)
            hook_instance.create_version.assert_called_with(
                'test-project', 'test-model', self.VERSION_INPUT)
    def testSuccessCreateVersion(self):
        with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
                as mock_hook:
            success_response = {'name': 'some-name', 'done': True}
            hook_instance = mock_hook.return_value
            hook_instance.create_version.return_value = success_response

            training_op = MLEngineVersionOperator(
                version=self.VERSION_INPUT,
                **self.VERSION_DEFAULT_ARGS)
            training_op.execute(None)

            mock_hook.assert_called_with(gcp_conn_id='google_cloud_default',
                                         delegate_to=None)
            # Make sure only 'create_version' is invoked on hook instance
            self.assertEqual(len(hook_instance.mock_calls), 1)
            hook_instance.create_version.assert_called_with(
                'test-project', 'test-model', self.VERSION_INPUT)
def deploy_tasks(model, parent_dag_name, child_dag_name, default_args,
                 PROJECT_ID, MODEL_NAME, MODEL_VERSION, MODEL_LOCATION):
    # Create inner dag
    dag = DAG("{0}.{1}".format(parent_dag_name, child_dag_name),
              default_args=default_args,
              schedule_interval=None)

    # Constants
    OTHER_VERSION_NAME = "v_{0}".format(
        datetime.datetime.now().strftime("%Y%m%d%H%M%S")[0:12])

    # Create model on ML-Engine
    bash_ml_engine_models_list_op = BashOperator(
        task_id="bash_ml_engine_models_list_{}_task".format(
            model.replace(".", "_")),
        xcom_push=True,
        bash_command="gcloud ml-engine models list --filter='name:{0}'".format(
            MODEL_NAME + model.replace(".", "_")),
        dag=dag)

    def check_if_model_already_exists(templates_dict, **kwargs):
        cur_model = templates_dict["model"].replace(".", "_")
        ml_engine_models_list = kwargs["ti"].xcom_pull(
            task_ids="bash_ml_engine_models_list_{}_task".format(cur_model))
        logging.info(
            "check_if_model_already_exists: {}: ml_engine_models_list = \n{}".
            format(cur_model, ml_engine_models_list))
        create_model_task = "ml_engine_create_model_{}_task".format(cur_model)
        dont_create_model_task = "dont_create_model_dummy_branch_{}_task".format(
            cur_model)
        if len(ml_engine_models_list
               ) == 0 or ml_engine_models_list == "Listed 0 items.":
            return create_model_task
        return dont_create_model_task

    check_if_model_already_exists_op = BranchPythonOperator(
        task_id="check_if_model_already_exists_{}_task".format(
            model.replace(".", "_")),
        templates_dict={"model": model.replace(".", "_")},
        python_callable=check_if_model_already_exists,
        provide_context=True,
        dag=dag)

    ml_engine_create_model_op = MLEngineModelOperator(
        task_id="ml_engine_create_model_{}_task".format(model.replace(
            ".", "_")),
        project_id=PROJECT_ID,
        model={"name": MODEL_NAME + model.replace(".", "_")},
        operation="create",
        dag=dag)

    create_model_dummy_op = DummyOperator(
        task_id="create_model_dummy_{}_task".format(model.replace(".", "_")),
        trigger_rule="all_done",
        dag=dag)

    dont_create_model_dummy_branch_op = DummyOperator(
        task_id="dont_create_model_dummy_branch_{}_task".format(
            model.replace(".", "_")),
        dag=dag)

    dont_create_model_dummy_op = DummyOperator(
        task_id="dont_create_model_dummy_{}_task".format(
            model.replace(".", "_")),
        trigger_rule="all_done",
        dag=dag)

    # Create version of model on ML-Engine
    bash_ml_engine_versions_list_op = BashOperator(
        task_id="bash_ml_engine_versions_list_{}_task".format(
            model.replace(".", "_")),
        xcom_push=True,
        bash_command=
        "gcloud ml-engine versions list --model {0} --filter='name:{1}'".
        format(MODEL_NAME + model.replace(".", "_"), MODEL_VERSION),
        dag=dag)

    def check_if_model_version_already_exists(templates_dict, **kwargs):
        cur_model = templates_dict["model"].replace(".", "_")
        ml_engine_versions_list = kwargs["ti"].xcom_pull(
            task_ids="bash_ml_engine_versions_list_{}_task".format(cur_model))
        logging.info(
            "check_if_model_version_already_exists: {}: ml_engine_versions_list = \n{}"
            .format(cur_model, ml_engine_versions_list))
        create_version_task = "ml_engine_create_version_{}_task".format(
            cur_model)
        create_other_version_task = "ml_engine_create_other_version_{}_task".format(
            cur_model)
        if len(ml_engine_versions_list
               ) == 0 or ml_engine_versions_list == "Listed 0 items.":
            return create_version_task
        return create_other_version_task

    check_if_model_version_already_exists_op = BranchPythonOperator(
        task_id="check_if_model_version_already_exists_{}_task".format(
            model.replace(".", "_")),
        templates_dict={"model": model.replace(".", "_")},
        python_callable=check_if_model_version_already_exists,
        provide_context=True,
        dag=dag)

    ml_engine_create_version_op = MLEngineVersionOperator(
        task_id="ml_engine_create_version_{}_task".format(
            model.replace(".", "_")),
        project_id=PROJECT_ID,
        model_name=MODEL_NAME + model.replace(".", "_"),
        version_name=MODEL_VERSION,
        version={
            "name": MODEL_VERSION,
            "deploymentUri": MODEL_LOCATION + model.replace(".", "_"),
            "runtimeVersion": "1.13",
            "framework": "TENSORFLOW",
            "pythonVersion": "3.5",
        },
        operation="create",
        dag=dag)

    ml_engine_create_other_version_op = MLEngineVersionOperator(
        task_id="ml_engine_create_other_version_{}_task".format(
            model.replace(".", "_")),
        project_id=PROJECT_ID,
        model_name=MODEL_NAME + model.replace(".", "_"),
        version_name=OTHER_VERSION_NAME,
        version={
            "name": OTHER_VERSION_NAME,
            "deploymentUri": MODEL_LOCATION + model.replace(".", "_"),
            "runtimeVersion": "1.13",
            "framework": "TENSORFLOW",
            "pythonVersion": "3.5",
        },
        operation="create",
        dag=dag)

    ml_engine_set_default_version_op = MLEngineVersionOperator(
        task_id="ml_engine_set_default_version_{}_task".format(
            model.replace(".", "_")),
        project_id=PROJECT_ID,
        model_name=MODEL_NAME + model.replace(".", "_"),
        version_name=MODEL_VERSION,
        version={"name": MODEL_VERSION},
        operation="set_default",
        dag=dag)

    ml_engine_set_default_other_version_op = MLEngineVersionOperator(
        task_id="ml_engine_set_default_other_version_{}_task".format(
            model.replace(".", "_")),
        project_id=PROJECT_ID,
        model_name=MODEL_NAME + model.replace(".", "_"),
        version_name=OTHER_VERSION_NAME,
        version={"name": OTHER_VERSION_NAME},
        operation="set_default",
        dag=dag)

    # Build dependency graph, set_upstream dependencies for all tasks
    check_if_model_already_exists_op.set_upstream(
        bash_ml_engine_models_list_op)

    ml_engine_create_model_op.set_upstream(check_if_model_already_exists_op)
    create_model_dummy_op.set_upstream(ml_engine_create_model_op)
    dont_create_model_dummy_branch_op.set_upstream(
        check_if_model_already_exists_op)
    dont_create_model_dummy_op.set_upstream(dont_create_model_dummy_branch_op)

    bash_ml_engine_versions_list_op.set_upstream(
        [dont_create_model_dummy_op, create_model_dummy_op])
    check_if_model_version_already_exists_op.set_upstream(
        bash_ml_engine_versions_list_op)

    ml_engine_create_version_op.set_upstream(
        check_if_model_version_already_exists_op)
    ml_engine_create_other_version_op.set_upstream(
        check_if_model_version_already_exists_op)

    ml_engine_set_default_version_op.set_upstream(ml_engine_create_version_op)
    ml_engine_set_default_other_version_op.set_upstream(
        ml_engine_create_other_version_op)

    return dag
        templates_dict={"model": model.replace(".", "_")},
        python_callable=check_if_model_version_already_exists,
        provide_context=True,
        dag=dag)

    OTHER_VERSION_NAME = "v_{0}".format(
        datetime.datetime.now().strftime("%Y%m%d%H%M%S")[0:12])

    ml_engine_create_version_op = MLEngineVersionOperator(
        task_id="ml_engine_create_version_{}_task".format(
            model.replace(".", "_")),
        project_id=PROJECT_ID,
        model_name=MODEL_NAME + model.replace(".", "_"),
        version_name=MODEL_VERSION,
        version={
            "name": MODEL_VERSION,
            "deploymentUri": MODEL_LOCATION + model.replace(".", "_"),
            "runtimeVersion": "1.13",
            "framework": "TENSORFLOW",
            "pythonVersion": "3.5",
        },
        operation="create",
        dag=dag)

    ml_engine_create_other_version_op = MLEngineVersionOperator(
        task_id="ml_engine_create_other_version_{}_task".format(
            model.replace(".", "_")),
        project_id=PROJECT_ID,
        model_name=MODEL_NAME + model.replace(".", "_"),
        version_name=OTHER_VERSION_NAME,
        version={
        provide_context=True,
        trigger_rule="none_failed",
        dag=dag
        )

    # MLEngineVersionOperator with operation set to "create" to create a new
    # version of our model
    ml_engine_create_version_op = MLEngineVersionOperator(
        task_id="ml_engine_create_version_{}_task"
                .format(model.replace(".", "_")),
        project_id=PROJECT_ID,
        model_name=MODEL_NAME,
        version_name=Variable.get("CURRENT_VERSION_NAME"),
        version={
            "name": Variable.get("CURRENT_VERSION_NAME"),
            "deploymentUri": MODEL_LOCATION + model.replace(".", "_"),
            "runtimeVersion": "2.1",
            "framework": "TENSORFLOW",
            "pythonVersion": "3.7",
            },
        operation="create",
        trigger_rule='none_failed',
        dag=dag
    )

    # MLEngineVersionOperator with operation set to "set_default" to set our
    # newly deployed version to be the default version.
    ml_engine_set_default_version_op = MLEngineVersionOperator(
        task_id="ml_engine_set_default_version_{}_task"
                .format(model.replace(".", "_")),
        project_id=PROJECT_ID,