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)