def test_create_execute(self, body, mock_hook): operator = GKECreateClusterOperator(project_id=TEST_GCP_PROJECT_ID, location=PROJECT_LOCATION, body=body, task_id=PROJECT_TASK_ID) operator.execute(None) mock_hook.return_value.create_cluster.assert_called_once_with( cluster=body, project_id=TEST_GCP_PROJECT_ID)
# [START composer_gkeoperator_affinity_airflow_1] # [START composer_gkeoperator_fullconfig_airflow_1] # TODO(developer): update with your values PROJECT_ID = "my-project-id" CLUSTER_ZONE = "us-west1-a" CLUSTER_NAME = "example-cluster" # [END composer_gkeoperator_minconfig_airflow_1] # [END composer_gkeoperator_templateconfig_airflow_1] # [END composer_gkeoperator_affinity_airflow_1] # [END composer_gkeoperator_fullconfig_airflow_1] CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1} # [END composer_gke_create_cluster_airflow_1] # [START composer_gke_create_cluster_airflow_1] create_cluster = GKECreateClusterOperator( task_id="create_cluster", project_id=PROJECT_ID, location=CLUSTER_ZONE, body=CLUSTER, ) # Using the BashOperator to create node pools is a workaround # In Airflow 2, because of https://github.com/apache/airflow/pull/17820 # Node pool creation can be done using the GKECreateClusterOperator create_node_pools = BashOperator( task_id="create_node_pools", bash_command=f"gcloud container node-pools create pool-0 \ --cluster {CLUSTER_NAME} \ --num-nodes 1 \ --zone {CLUSTER_ZONE} \ && gcloud container node-pools create pool-1 \ --cluster {CLUSTER_NAME} \ --num-nodes 1 \
def test_create_execute_error_body(self, body, mock_hook): with self.assertRaises(AirflowException): GKECreateClusterOperator(project_id=TEST_GCP_PROJECT_ID, location=PROJECT_LOCATION, body=body, task_id=PROJECT_TASK_ID)
def test_create_execute_error_project_id(self, mock_hook): with self.assertRaises(AirflowException): GKECreateClusterOperator(location=PROJECT_LOCATION, body=PROJECT_BODY, task_id=PROJECT_TASK_ID)
def test_create_execute_error_location(self, mock_hook): with pytest.raises(AirflowException): GKECreateClusterOperator(project_id=TEST_GCP_PROJECT_ID, body=PROJECT_BODY, task_id=PROJECT_TASK_ID)
GCP_LOCATION = os.environ.get("GCP_GKE_LOCATION", "europe-north1-a") CLUSTER_NAME = os.environ.get("GCP_GKE_CLUSTER_NAME", "cluster-name") CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1} default_args = {"start_date": days_ago(1)} with models.DAG( "example_gcp_gke", default_args=default_args, schedule_interval=None, # Override to match your needs tags=['example'], ) as dag: create_cluster = GKECreateClusterOperator( task_id="create_cluster", project_id=GCP_PROJECT_ID, location=GCP_LOCATION, body=CLUSTER, ) pod_task = GKEStartPodOperator( task_id="pod_task", project_id=GCP_PROJECT_ID, location=GCP_LOCATION, cluster_name=CLUSTER_NAME, namespace="default", image="perl", name="test-pod", ) pod_task_xcom = GKEStartPodOperator( task_id="pod_task_xcom",