def test_get_namespace(self):
kubernetes_hook_with_namespace = KubernetesHook(
conn_id='kubernetes_with_namespace')
kubernetes_hook_without_namespace = KubernetesHook(
conn_id='kubernetes_default_kube_config')
assert kubernetes_hook_with_namespace.get_namespace(
) == 'mock_namespace'
assert kubernetes_hook_without_namespace.get_namespace() == 'default'
def test_get_namespace(self):
kubernetes_hook_with_namespace = KubernetesHook(
conn_id='kubernetes_with_namespace')
kubernetes_hook_without_namespace = KubernetesHook(
conn_id='kubernetes_default_kube_config')
self.assertEqual(kubernetes_hook_with_namespace.get_namespace(),
'mock_namespace')
self.assertEqual(kubernetes_hook_without_namespace.get_namespace(),
'default')
def test_kube_config_path(self, mock_kube_config_loader,
mock_kube_config_merger):
kubernetes_hook = KubernetesHook(conn_id='kubernetes_kube_config_path')
api_conn = kubernetes_hook.get_conn()
mock_kube_config_loader.assert_called_once_with("path/to/file")
mock_kube_config_merger.assert_called_once()
self.assertIsInstance(api_conn, kubernetes.client.api_client.ApiClient)
def test_kube_config_connection(self, mock_kube_config_loader,
mock_kube_config_merger, mock_tempfile):
kubernetes_hook = KubernetesHook(conn_id='kubernetes_kube_config')
api_conn = kubernetes_hook.get_conn()
mock_tempfile.is_called_once()
mock_kube_config_loader.assert_called_once()
mock_kube_config_merger.assert_called_once()
self.assertIsInstance(api_conn, kubernetes.client.api_client.ApiClient)
def test_should_raise_exception_on_invalid_configuration(self, conn_uri):
with mock.patch.dict("os.environ",
AIRFLOW_CONN_KUBERNETES_DEFAULT=conn_uri
), self.assertRaisesRegex(
AirflowException,
"Invalid connection configuration"):
kubernetes_hook = KubernetesHook()
kubernetes_hook.get_conn()
def get_hook(self):
hook = KubernetesHook(
conn_id=self.kubernetes_conn_id,
in_cluster=self.in_cluster,
config_file=self.config_file,
cluster_context=self.cluster_context,
)
self._patch_deprecated_k8s_settings(hook)
return hook
def execute(self, context):
self.log.info("Creating sparkApplication")
hook = KubernetesHook(conn_id=self.kubernetes_conn_id)
response = hook.create_custom_resource_definition(
group="sparkoperator.k8s.io",
version="v1beta2",
plural="sparkapplications",
body=self.application_file,
namespace=self.namespace)
return response
def execute(self, context):
self.log.info("Creating sparkApplication")
hook = KubernetesHook(conn_id=self.kubernetes_conn_id)
response = hook.create_custom_object(
group=self.api_group,
version=self.api_version,
plural="sparkapplications",
body=self.application_file,
namespace=self.namespace,
)
return response
def test_default_kube_config_connection(
self,
mock_kube_config_loader,
mock_kube_config_merger,
):
kubernetes_hook = KubernetesHook(
conn_id='kubernetes_default_kube_config')
api_conn = kubernetes_hook.get_conn()
mock_kube_config_loader.assert_called_once_with("/mock/config")
mock_kube_config_merger.assert_called_once()
self.assertIsInstance(api_conn, kubernetes.client.api_client.ApiClient)
def __init__(
self,
*,
application_name: str,
attach_log: bool = False,
namespace: Optional[str] = None,
kubernetes_conn_id: str = "kubernetes_default",
**kwargs,
) -> None:
super().__init__(**kwargs)
self.application_name = application_name
self.attach_log = attach_log
self.namespace = namespace
self.kubernetes_conn_id = kubernetes_conn_id
self.hook = KubernetesHook(conn_id=self.kubernetes_conn_id)
def __init__(
self,
*,
application_name: str,
attach_log: bool = False,
namespace: Optional[str] = None,
kubernetes_conn_id: str = "kubernetes_default",
api_group: str = 'sparkoperator.k8s.io',
api_version: str = 'v1beta2',
**kwargs,
) -> None:
super().__init__(**kwargs)
self.application_name = application_name
self.attach_log = attach_log
self.namespace = namespace
self.kubernetes_conn_id = kubernetes_conn_id
self.hook = KubernetesHook(conn_id=self.kubernetes_conn_id)
self.api_group = api_group
self.api_version = api_version
def poke(self, context: Dict):
self.log.info("Poking: %s", self.application_name)
hook = KubernetesHook(conn_id=self.kubernetes_conn_id)
response = hook.get_custom_resource_definition(
group="sparkoperator.k8s.io",
version="v1beta2",
plural="sparkapplications",
name=self.application_name,
namespace=self.namespace)
try:
application_state = response['status']['applicationState']['state']
except KeyError:
return False
if application_state in self.FAILURE_STATES:
raise AirflowException("Spark application failed with state: %s" % application_state)
elif application_state in self.SUCCESS_STATES:
self.log.info("Spark application ended successfully")
return True
else:
self.log.info("Spark application is still in state: %s", application_state)
return False
def test_in_cluster_connection(self, mock_kube_config_loader):
kubernetes_hook = KubernetesHook(conn_id='kubernetes_in_cluster')
api_conn = kubernetes_hook.get_conn()
mock_kube_config_loader.assert_called_once()
self.assertIsInstance(api_conn, kubernetes.client.api_client.ApiClient)