def test_get_hook_defined(
self,
_: mock.MagicMock, # force creation of _get_hook
sql_sensor_poke_mock: mock.MagicMock,
):
""" Test for case when ``SqlSensor._get_hook`` method is present. """
op = ClickHouseSqlSensor(task_id='_', sql='')
context = dict(some_test_context=True)
op.poke(context)
sql_sensor_poke_mock.assert_called_once_with(context)
def test_clickhouse_sql_sensor_poke_failure(self):
op = ClickHouseSqlSensor(task_id='_', sql='', failure=[1].__contains__)
self._get_records_mock.return_value = []
with self.subTest(self._get_records_mock.return_value):
self.assertFalse(op.poke(None))
self._get_records_mock.return_value = [[1]]
with self.subTest(self._get_records_mock.return_value):
with self.assertRaises(AirflowException):
op.poke(context=None)
def test_clickhouse_sql_sensor_poke_failure_success(self):
op = ClickHouseSqlSensor(
task_id='_', sql='',
failure=[1].__contains__, success=[2].__contains__,
)
for return_value, expected_result in (
([], False),
([[2]], True),
):
with self.subTest(return_value):
self._get_records_mock.return_value = return_value
self.assertEqual(expected_result, op.poke(context=None))
self._get_records_mock.return_value = [[1]]
with self.subTest(self._get_records_mock.return_value):
with self.assertRaises(AirflowException):
op.poke(context=None)
def test_clickhouse_sql_sensor_poke_success(self):
op = ClickHouseSqlSensor(task_id='_', sql='', success=[1].__contains__)
for return_value, expected_result in (
([], False),
([[1]], True),
([['1']], False),
):
with self.subTest(return_value):
self._get_records_mock.return_value = return_value
self.assertEqual(expected_result, op.poke(context=None))
def test_clickhouse_sql_sensor_poke_invalid_success(self):
op = ClickHouseSqlSensor(task_id='_', sql='', success=[1])
self._get_records_mock.return_value = [[1]]
with self.assertRaises(AirflowException):
op.poke(context=None)
def test_clickhouse_sql_sensor_poke_fail_on_empty(self):
op = ClickHouseSqlSensor(task_id='_', sql='', fail_on_empty=True)
self._get_records_mock.return_value = []
with self.assertRaises(AirflowException):
op.poke(context=None)