def test_get_upstream_partitions_for_partition_range_monthly_downstream_daily_upstream():
downstream_partitions_def = MonthlyPartitionsDefinition(start_date="2021-05-01")
upstream_partitions_def = DailyPartitionsDefinition(start_date="2021-05-01")
result = TimeWindowPartitionMapping().get_upstream_partitions_for_partition_range(
PartitionKeyRange("2021-05-01", "2021-07-01"),
downstream_partitions_def,
upstream_partitions_def,
)
assert result == PartitionKeyRange("2021-05-01", "2021-07-31")
def test_get_upstream_partitions_for_partition_range_same_partitioning_different_formats():
downstream_partitions_def = DailyPartitionsDefinition(start_date="2021-05-05")
upstream_partitions_def = DailyPartitionsDefinition(start_date="2021/05/05", fmt="%Y/%m/%d")
result = TimeWindowPartitionMapping().get_upstream_partitions_for_partition_range(
PartitionKeyRange("2021-05-07", "2021-05-09"),
downstream_partitions_def,
upstream_partitions_def,
)
assert result == PartitionKeyRange("2021/05/07", "2021/05/09")
def test_get_upstream_partitions_for_partition_range_daily_downstream_hourly_upstream():
downstream_partitions_def = DailyPartitionsDefinition(start_date="2021-05-05")
upstream_partitions_def = HourlyPartitionsDefinition(start_date="2021-05-05-00:00")
result = TimeWindowPartitionMapping().get_upstream_partitions_for_partition_range(
PartitionKeyRange("2021-05-07", "2021-05-07"),
downstream_partitions_def,
upstream_partitions_def,
)
assert result == PartitionKeyRange("2021-05-07-00:00", "2021-05-07-23:00")
result = TimeWindowPartitionMapping().get_upstream_partitions_for_partition_range(
PartitionKeyRange("2021-05-07", "2021-05-09"),
downstream_partitions_def,
upstream_partitions_def,
)
assert result == PartitionKeyRange("2021-05-07-00:00", "2021-05-09-23:00")
def test_get_upstream_partitions_for_partition_range_same_partitioning():
downstream_partitions_def = DailyPartitionsDefinition(start_date="2021-05-05")
upstream_partitions_def = DailyPartitionsDefinition(start_date="2021-05-05")
# single partition key
partition_key_range = PartitionKeyRange("2021-05-07", "2021-05-07")
result = TimeWindowPartitionMapping().get_upstream_partitions_for_partition_range(
partition_key_range,
downstream_partitions_def,
upstream_partitions_def,
)
assert partition_key_range == result
# range of partition keys
partition_key_range = PartitionKeyRange("2021-05-07", "2021-05-09")
result = TimeWindowPartitionMapping().get_upstream_partitions_for_partition_range(
partition_key_range,
downstream_partitions_def,
upstream_partitions_def,
)
assert partition_key_range == result