From ea078ffc6679cbed3cc7b2a8466c999a12f1c1a1 Mon Sep 17 00:00:00 2001
From: Jeroen Dries <jeroen.dries@vito.be>
Date: Fri, 7 Jun 2024 16:24:50 +0200
Subject: [PATCH] do not push down resample_cube_spatial in case an 'apply'
 process is in between
 https://github.com/Open-EO/openeo-geopyspark-driver/issues/790

---
 openeo_driver/dry_run.py |  6 +++++-
 tests/test_dry_run.py    | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/openeo_driver/dry_run.py b/openeo_driver/dry_run.py
index c1bf71fb..ef3bb5b6 100644
--- a/openeo_driver/dry_run.py
+++ b/openeo_driver/dry_run.py
@@ -706,6 +706,10 @@ def apply_tiles_spatiotemporal(self, process, context: Optional[dict] = None) ->
         else:
             return self
 
+    def apply(self, process: dict, *, context: Optional[dict] = None, env: EvalEnv) -> "DriverDataCube":
+        cube = self._process("apply", {})
+        return cube
+
     def apply_neighborhood(
         self, process, *, size: List[dict], overlap: List[dict], context: Optional[dict] = None, env: EvalEnv
     ) -> "DriverDataCube":
@@ -771,7 +775,7 @@ def fit_class_random_forest(self, predictors, target, training, num_trees, mtry)
 
     # TODO: some methods need metadata manipulation?
 
-    apply = _nop
+
     apply_tiles = _nop
 
     reduce = _nop
diff --git a/tests/test_dry_run.py b/tests/test_dry_run.py
index fe9d2720..8611d3da 100644
--- a/tests/test_dry_run.py
+++ b/tests/test_dry_run.py
@@ -988,6 +988,25 @@ def test_resample_filter_spatial(dry_run_env, dry_run_tracer):
     }
 
 
+
+def test_resample_cube_spatial(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"), connection=None)
+    cube = cube.apply(lambda x: x + 1)
+    cube = cube.resample_cube_spatial(target=DataCube(PGNode("load_collection", id="S2_FOOBAR"), connection=None),method="max")
+
+    pg = cube.flat_graph()
+    res = evaluate(pg, env=dry_run_env,do_dry_run=False)
+
+    source_constraints = dry_run_tracer.get_source_constraints(merge=True)
+    assert len(source_constraints) == 2
+    src, constraints = source_constraints[0]
+    assert src == ("load_collection", ("S2_FOOBAR", ()))
+
+    assert constraints == { 'process_type': [ ProcessType.FOCAL_SPACE]}
+
+
+
 def test_auto_align(dry_run_env, dry_run_tracer):
     polygon = {"type": "Polygon", "coordinates": [[(0.1, 0.1), (3, 5), (8, 2), (0.1, 0.1)]]}
     cube = DataCube(PGNode("load_collection", id="ESA_WORLDCOVER_10M_2020_V1"), connection=None)