def test_mask_polygon_and_load_collection_spatial_extent(dry_run_env, dry_run_tracer):
polygon = {"type": "Polygon", "coordinates": [[(0, 0), (3, 5), (8, 2), (0, 0)]]}
cube = DataCube(PGNode(
"load_collection", id="S2_FOOBAR",
spatial_extent={"west": -1, "south": -1, "east": 10, "north": 10}
), connection=None)
cube = cube.mask_polygon(mask=polygon)
pg = cube.flat_graph()
res = evaluate(pg, env=dry_run_env)
source_constraints = dry_run_tracer.get_source_constraints(merge=True)
assert len(source_constraints) == 1
src, constraints = source_constraints[0]
assert src == ("load_collection", ("S2_FOOBAR", ()))
assert constraints == {
"spatial_extent": {"west": -1, "south": -1, "east": 10, "north": 10, "crs": "EPSG:4326"}
}
def test_mask_polygon_only(dry_run_env, dry_run_tracer, inside, replacement, expect_spatial_extent):
polygon = {"type": "Polygon", "coordinates": [[(0, 0), (3, 5), (8, 2), (0, 0)]]}
cube = DataCube(PGNode("load_collection", id="S2_FOOBAR"), connection=None)
cube = cube.mask_polygon(mask=polygon, inside=inside, replacement=replacement)
pg = cube.flat_graph()
res = evaluate(pg, env=dry_run_env)
source_constraints = dry_run_tracer.get_source_constraints(merge=True)
assert len(source_constraints) == 1
src, constraints = source_constraints[0]
assert src == ("load_collection", ("S2_FOOBAR", ()))
if expect_spatial_extent:
expected = {
"spatial_extent": {"west": 0.0, "south": 0.0, "east": 8.0, "north": 5.0, "crs": "EPSG:4326"}
}
else:
expected = {}
assert constraints == expected
def test_mask_polygon_and_filter_bbox(dry_run_env, dry_run_tracer, bbox_first):
polygon = {"type": "Polygon", "coordinates": [[(0, 0), (3, 5), (8, 2), (0, 0)]]}
bbox = {"west": -1, "south": -1, "east": 9, "north": 9, "crs": "EPSG:4326"}
# Use client lib to build process graph in flexible way
cube = DataCube(PGNode("load_collection", id="S2_FOOBAR"), connection=None)
if bbox_first:
cube = cube.filter_bbox(bbox=bbox).mask_polygon(mask=polygon)
else:
cube = cube.mask_polygon(mask=polygon).filter_bbox(bbox=bbox)
pg = cube.flat_graph()
res = evaluate(pg, env=dry_run_env)
source_constraints = dry_run_tracer.get_source_constraints(merge=True)
assert len(source_constraints) == 1
src, constraints = source_constraints[0]
assert src == ("load_collection", ("S2_FOOBAR", ()))
assert constraints == {
"spatial_extent": {"west": -1, "south": -1, "east": 9, "north": 9, "crs": "EPSG:4326"}
}