def test_sql_sensor_postgres_poke(self, mock_hook): t = SqlSensor( task_id='sql_sensor_check', conn_id='postgres_default', sql="SELECT 1", ) mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records mock_get_records.return_value = [] self.assertFalse(t.poke(None)) mock_get_records.return_value = [['1']] self.assertTrue(t.poke(None))
def test_sql_sensor_postgres_poke_success(self, mock_hook): t = SqlSensor(task_id='sql_sensor_check', conn_id='postgres_default', sql="SELECT 1", success=lambda x: x in [1]) mock_hook.get_connection('postgres_default').conn_type = "postgres" mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records mock_get_records.return_value = [] self.assertFalse(t.poke(None)) mock_get_records.return_value = [[1]] self.assertTrue(t.poke(None)) mock_get_records.return_value = [['1']] self.assertFalse(t.poke(None))
def test_sql_sensor_postgres_poke(self, mock_hook): op = SqlSensor( task_id='sql_sensor_check', conn_id='postgres_default', sql="SELECT 1", ) mock_hook.get_connection('postgres_default').conn_type = "postgres" mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records mock_get_records.return_value = [] self.assertFalse(op.poke(None)) mock_get_records.return_value = [[None]] self.assertFalse(op.poke(None)) mock_get_records.return_value = [['None']] self.assertTrue(op.poke(None)) mock_get_records.return_value = [[0.0]] self.assertFalse(op.poke(None)) mock_get_records.return_value = [[0]] self.assertFalse(op.poke(None)) mock_get_records.return_value = [['0']] self.assertTrue(op.poke(None)) mock_get_records.return_value = [['1']] self.assertTrue(op.poke(None))
def test_sql_sensor_postgres_poke_failure_success(self, mock_hook): op = SqlSensor(task_id='sql_sensor_check', conn_id='postgres_default', sql="SELECT 1", failure=lambda x: x in [1], success=lambda x: x in [2]) mock_hook.get_connection('postgres_default').conn_type = "postgres" mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records mock_get_records.return_value = [] self.assertFalse(op.poke(None)) mock_get_records.return_value = [[1]] self.assertRaises(AirflowException, op.poke, None) mock_get_records.return_value = [[2]] self.assertTrue(op.poke(None))
def test_sql_sensor_postgres_check_allow_null(self, mock_hook): t1 = SqlSensor(task_id='sql_sensor_check', conn_id='postgres_default', sql="SELECT NULL", allow_null=True) mock_hook.get_connection('postgres_default').conn_type = "postgres" mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records mock_get_records.return_value = [[None]] self.assertTrue(t1.poke(None)) t2 = SqlSensor(task_id='sql_sensor_check', conn_id='postgres_default', sql="SELECT NULL", allow_null=False) mock_hook.get_connection('postgres_default').conn_type = "postgres" mock_get_records = mock_hook.get_connection.return_value.get_hook.return_value.get_records mock_get_records.return_value = [[None]] self.assertFalse(t2.poke(None))