TASK_ID = "test-automl-hook" GCP_PROJECT_ID = "test-project" GCP_LOCATION = "test-location" MODEL_NAME = "test_model" MODEL_ID = "projects/198907790164/locations/us-central1/models/TBL9195602771183665152" DATASET_ID = "TBL123456789" MODEL = { "display_name": MODEL_NAME, "dataset_id": DATASET_ID, "tables_model_metadata": { "train_budget_milli_node_hours": 1000 }, } LOCATION_PATH = AutoMlClient.location_path(GCP_PROJECT_ID, GCP_LOCATION) MODEL_PATH = PredictionServiceClient.model_path(GCP_PROJECT_ID, GCP_LOCATION, MODEL_ID) INPUT_CONFIG = {"input": "value"} OUTPUT_CONFIG = {"output": "value"} PAYLOAD = {"test": "payload"} DATASET = {"dataset_id": "data"} MASK = {"field": "mask"} class TestAutoMLTrainModelOperator(unittest.TestCase): @mock.patch( "airflow.gcp.operators.automl.AutoMLTrainModelOperator.xcom_push") @mock.patch("airflow.gcp.hooks.automl.CloudAutoMLHook.create_model") @mock.patch( "airflow.gcp.hooks.automl.CloudAutoMLHook.extract_object_id", return_value=MODEL_ID,
def __init__(self, project, location, automl_model_name): # The default client behaviour excepts a global model, set the specific EU endpoint for EU models self.client = PredictionServiceClient(client_options={"api_endpoint": "eu-automl.googleapis.com:443"}) self.path = PredictionServiceClient.model_path(project, location, automl_model_name)