def test_driver_logging_error(self, mock_log_call, warn_log_call,
mock_get_namespaced_crd, mock_kube_conn):
sensor = SparkKubernetesSensor(
application_name="spark_pi",
attach_log=True,
dag=self.dag,
task_id="test_task_id",
)
sensor.poke(None)
warn_log_call.assert_called_once()
def test_unknown_application(self, mock_get_namespaced_crd, mock_kubernetes_hook):
sensor = SparkKubernetesSensor(application_name="spark_pi", dag=self.dag, task_id="test_task_id")
with pytest.raises(AirflowException):
sensor.poke(None)
mock_kubernetes_hook.assert_called_once_with()
mock_get_namespaced_crd.assert_called_once_with(
group="sparkoperator.k8s.io",
name="spark_pi",
namespace="default",
plural="sparkapplications",
version="v1beta2",
)
def test_driver_logging_completed(self, mock_log_call, info_log_call,
mock_get_namespaced_crd, mock_kube_conn):
sensor = SparkKubernetesSensor(
application_name="spark_pi",
attach_log=True,
dag=self.dag,
task_id="test_task_id",
)
sensor.poke(None)
mock_log_call.assert_called_once_with("spark_pi-driver")
log_info_call = info_log_call.mock_calls[1]
log_value = log_info_call[1][0]
self.assertEqual(log_value, TEST_POD_LOG_RESULT)
def test_driver_logging_failure(self, mock_log_call, error_log_call,
mock_get_namespaced_crd, mock_kube_conn):
sensor = SparkKubernetesSensor(
application_name="spark_pi",
attach_log=True,
dag=self.dag,
task_id="test_task_id",
)
with pytest.raises(AirflowException):
sensor.poke(None)
mock_log_call.assert_called_once_with("spark-pi-driver",
namespace="default")
error_log_call.assert_called_once_with(TEST_POD_LOG_RESULT)
def test_namespace_from_connection(self, mock_get_namespaced_crd,
mock_kubernetes_hook):
sensor = SparkKubernetesSensor(
application_name='spark_pi',
dag=self.dag,
kubernetes_conn_id='kubernetes_with_namespace',
task_id='test_task_id')
sensor.poke(None)
mock_kubernetes_hook.assert_called_once_with()
mock_get_namespaced_crd.assert_called_once_with(
group='sparkoperator.k8s.io',
name='spark_pi',
namespace='mock_namespace',
plural='sparkapplications',
version='v1beta2')
def test_namespace_from_connection(self, mock_get_namespaced_crd, mock_kubernetes_hook):
sensor = SparkKubernetesSensor(
application_name="spark_pi",
dag=self.dag,
kubernetes_conn_id="kubernetes_with_namespace",
task_id="test_task_id",
)
sensor.poke(None)
mock_kubernetes_hook.assert_called_once_with()
mock_get_namespaced_crd.assert_called_once_with(
group="sparkoperator.k8s.io",
name="spark_pi",
namespace="mock_namespace",
plural="sparkapplications",
version="v1beta2",
)
def test_pending_rerun_application(self, mock_get_namespaced_crd, mock_kubernetes_hook):
sensor = SparkKubernetesSensor(application_name="spark_pi", dag=self.dag, task_id="test_task_id")
self.assertFalse(sensor.poke(None))
mock_kubernetes_hook.assert_called_once_with()
mock_get_namespaced_crd.assert_called_once_with(
group="sparkoperator.k8s.io",
name="spark_pi",
namespace="default",
plural="sparkapplications",
version="v1beta2",
)
def test_api_group_and_version_from_sensor(self, mock_get_namespaced_crd,
mock_kubernetes_hook):
api_group = 'sparkoperator.example.com'
api_version = 'v1alpha1'
sensor = SparkKubernetesSensor(
application_name="spark_pi",
dag=self.dag,
kubernetes_conn_id="kubernetes_with_namespace",
task_id="test_task_id",
api_group=api_group,
api_version=api_version,
)
sensor.poke(None)
mock_kubernetes_hook.assert_called_once_with()
mock_get_namespaced_crd.assert_called_once_with(
group=api_group,
name="spark_pi",
namespace="mock_namespace",
plural="sparkapplications",
version=api_version,
)
def test_pending_rerun_application(self, mock_get_namespaced_crd,
mock_kubernetes_hook):
sensor = SparkKubernetesSensor(application_name='spark_pi',
dag=self.dag,
task_id='test_task_id')
self.assertFalse(sensor.poke(None))
mock_kubernetes_hook.assert_called_once_with()
mock_get_namespaced_crd.assert_called_once_with(
group='sparkoperator.k8s.io',
name='spark_pi',
namespace='default',
plural='sparkapplications',
version='v1beta2')
