def test_poke_nondefault_args(self, mock_check_for_partition):
table_name = 'my_table'
expression = 'col=val'
aws_conn_id = 'my_aws_conn_id'
region_name = 'us-west-2'
database_name = 'my_db'
poke_interval = 2
timeout = 3
op = AwsGlueCatalogPartitionSensor(
task_id=self.task_id,
table_name=table_name,
expression=expression,
aws_conn_id=aws_conn_id,
region_name=region_name,
database_name=database_name,
poke_interval=poke_interval,
timeout=timeout,
)
op.poke(None)
self.assertEqual(op.hook.region_name, region_name)
self.assertEqual(op.hook.aws_conn_id, aws_conn_id)
self.assertEqual(op.poke_interval, poke_interval)
self.assertEqual(op.timeout, timeout)
mock_check_for_partition.assert_called_once_with(database_name, table_name, expression)
def test_poke_default_args(self, mock_check_for_partition):
table_name = 'test_glue_catalog_partition_sensor_tbl'
op = AwsGlueCatalogPartitionSensor(task_id=self.task_id, table_name=table_name)
op.poke(None)
self.assertEqual(op.hook.region_name, None)
self.assertEqual(op.hook.aws_conn_id, 'aws_default')
mock_check_for_partition.assert_called_once_with('default', table_name, "ds='{{ ds }}'")
def test_dot_notation(self, mock_check_for_partition):
db_table = 'my_db.my_tbl'
op = AwsGlueCatalogPartitionSensor(task_id=self.task_id,
table_name=db_table)
op.poke(None)
mock_check_for_partition.assert_called_once_with(
'my_db', 'my_tbl', "ds='{{ ds }}'")
def test_poke_false(self, mock_check_for_partition):
mock_check_for_partition.return_value = False
op = AwsGlueCatalogPartitionSensor(task_id=self.task_id, table_name='tbl')
self.assertFalse(op.poke(None))