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(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))