def test_create_execute_error_location(self, mock_hook):
with self.assertRaises(AirflowException):
operator = GKEClusterCreateOperator(project_id=TEST_GCP_PROJECT_ID,
body=PROJECT_BODY,
task_id=PROJECT_TASK_ID)
operator.execute(None)
mock_hook.return_value.create_cluster.assert_not_called()
def test_create_execute(self, body, mock_hook):
operator = GKEClusterCreateOperator(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)
def test_create_execute_error_body(self, body, mock_hook):
with self.assertRaises(AirflowException):
GKEClusterCreateOperator(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):
GKEClusterCreateOperator(location=PROJECT_LOCATION,
body=PROJECT_BODY,
task_id=PROJECT_TASK_ID)
GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
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
) as dag:
create_cluster = GKEClusterCreateOperator(
task_id="create_cluster",
project_id=GCP_PROJECT_ID,
location=GCP_LOCATION,
body=CLUSTER,
)
pod_task = GKEPodOperator(
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 = GKEPodOperator(
task_id="pod_task_xcom",
