def setUp(self):
self.gke_op = GKEStartPodOperator(project_id=TEST_GCP_PROJECT_ID,
location=PROJECT_LOCATION,
cluster_name=CLUSTER_NAME,
task_id=PROJECT_TASK_ID,
name=TASK_NAME,
namespace=NAMESPACE,
image=IMAGE)
class TestGKEPodOperator(unittest.TestCase):
def setUp(self):
self.gke_op = GKEStartPodOperator(project_id=TEST_GCP_PROJECT_ID,
location=PROJECT_LOCATION,
cluster_name=CLUSTER_NAME,
task_id=PROJECT_TASK_ID,
name=TASK_NAME,
namespace=NAMESPACE,
image=IMAGE)
def test_template_fields(self):
self.assertTrue(
set(KubernetesPodOperator.template_fields).issubset(
GKEStartPodOperator.template_fields))
# pylint: disable=unused-argument
@mock.patch("airflow.hooks.base_hook.BaseHook.get_connections",
return_value=[Connection(extra=json.dumps({}))])
@mock.patch(
'airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute'
)
@mock.patch('tempfile.NamedTemporaryFile')
@mock.patch("subprocess.check_call")
@mock.patch.dict(os.environ, {CREDENTIALS: '/tmp/local-creds'})
def test_execute_conn_id_none(self, proc_mock, file_mock, exec_mock,
get_conn):
type(
file_mock.return_value.__enter__.return_value).name = PropertyMock(
side_effect=[FILE_NAME])
def assert_credentials(*args, **kwargs):
# since we passed in keyfile_path we should get a file
self.assertIn(CREDENTIALS, os.environ)
self.assertEqual(os.environ[CREDENTIALS], '/tmp/local-creds')
proc_mock.side_effect = assert_credentials
self.gke_op.execute(None)
# Assert Environment Variable is being set correctly
self.assertIn(KUBE_ENV_VAR, os.environ)
self.assertEqual(os.environ[KUBE_ENV_VAR], FILE_NAME)
# Assert the gcloud command being called correctly
proc_mock.assert_called_once_with(
GCLOUD_COMMAND.format(CLUSTER_NAME, PROJECT_LOCATION,
TEST_GCP_PROJECT_ID).split())
self.assertEqual(self.gke_op.config_file, FILE_NAME)
# pylint: disable=unused-argument
@mock.patch(
"airflow.hooks.base_hook.BaseHook.get_connections",
return_value=[
Connection(extra=json.dumps(
{'extra__google_cloud_platform__key_path': '/path/to/file'}))
])
@mock.patch(
'airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute'
)
@mock.patch('tempfile.NamedTemporaryFile')
@mock.patch("subprocess.check_call")
@mock.patch.dict(os.environ, {})
def test_execute_conn_id_path(self, proc_mock, file_mock, exec_mock,
get_con_mock):
type(
file_mock.return_value.__enter__.return_value).name = PropertyMock(
side_effect=[FILE_NAME])
def assert_credentials(*args, **kwargs):
# since we passed in keyfile_path we should get a file
self.assertIn(CREDENTIALS, os.environ)
self.assertEqual(os.environ[CREDENTIALS], '/path/to/file')
proc_mock.side_effect = assert_credentials
self.gke_op.execute(None)
# Assert Environment Variable is being set correctly
self.assertIn(KUBE_ENV_VAR, os.environ)
self.assertEqual(os.environ[KUBE_ENV_VAR], FILE_NAME)
# Assert the gcloud command being called correctly
proc_mock.assert_called_once_with(
GCLOUD_COMMAND.format(CLUSTER_NAME, PROJECT_LOCATION,
TEST_GCP_PROJECT_ID).split())
self.assertEqual(self.gke_op.config_file, FILE_NAME)
# pylint: disable=unused-argument
@mock.patch.dict(os.environ, {})
@mock.patch("airflow.hooks.base_hook.BaseHook.get_connections",
return_value=[
Connection(extra=json.dumps({
"extra__google_cloud_platform__keyfile_dict":
'{"private_key": "r4nd0m_k3y"}'
}))
])
@mock.patch(
'airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute'
)
@mock.patch('tempfile.NamedTemporaryFile')
@mock.patch("subprocess.check_call")
def test_execute_conn_id_dict(self, proc_mock, file_mock, exec_mock,
get_con_mock):
type(
file_mock.return_value.__enter__.return_value).name = PropertyMock(
side_effect=[FILE_NAME, '/path/to/new-file'])
def assert_credentials(*args, **kwargs):
# since we passed in keyfile_dict we should get a new file
self.assertIn(CREDENTIALS, os.environ)
self.assertEqual(os.environ[CREDENTIALS], '/path/to/new-file')
proc_mock.side_effect = assert_credentials
self.gke_op.execute(None)
# Assert Environment Variable is being set correctly
self.assertIn(KUBE_ENV_VAR, os.environ)
self.assertEqual(os.environ[KUBE_ENV_VAR], FILE_NAME)
# Assert the gcloud command being called correctly
proc_mock.assert_called_once_with(
GCLOUD_COMMAND.format(CLUSTER_NAME, PROJECT_LOCATION,
TEST_GCP_PROJECT_ID).split())
self.assertEqual(self.gke_op.config_file, FILE_NAME)
"example_gcp_gke",
default_args=default_args,
schedule_interval=None, # Override to match your needs
) 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",
project_id=GCP_PROJECT_ID,
location=GCP_LOCATION,
cluster_name=CLUSTER_NAME,
do_xcom_push=True,
namespace="default",
image="alpine",
cmds=[
"sh", "-c",