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