Issue #391/#651 more review tweaks

Author: soxofaan

openeo/internal/graph_building.py

@@ -458,6 +458,8 @@ class MultiLeafGraph(FlatGraphableMixin):
     Container for process graphs with multiple leaf/result nodes.
     """
 
+    __slots__ = ["_leaves"]
+
     def __init__(self, leaves: Iterable[FlatGraphableMixin]):
         self._leaves = list(leaves)
 

openeo/rest/multiresult.py

@@ -34,6 +34,8 @@ class MultiResult(FlatGraphableMixin):
     .. versionadded:: 0.35.0
     """
 
+    __slots__ = ("_multi_leaf_graph", "_connection")
+
     def __init__(self, leaves: List[FlatGraphableMixin], connection: Optional[Connection] = None):
         """
         Build a :py:class:`MultiResult` instance from multiple leaf nodes
@@ -47,11 +49,14 @@ def __init__(self, leaves: List[FlatGraphableMixin], connection: Optional[Connec
             are not already associated with a connection.
         """
         self._multi_leaf_graph = MultiLeafGraph(leaves=leaves)
-        self._connection = self._common_connection(leaves=leaves, connection=connection)
+        self._connection = self._extract_connection(leaves=leaves, connection=connection)
 
     @staticmethod
-    def _common_connection(leaves: List[FlatGraphableMixin], connection: Optional[Connection] = None) -> Connection:
-        """Find common connection. Fails if there are multiple or none."""
+    def _extract_connection(leaves: List[FlatGraphableMixin], connection: Optional[Connection] = None) -> Connection:
+        """
+        Extract common connection from leaves and/or explicitly provided connection.
+        Fails if there are multiple or none.
+        """
         connections = set()
         if connection:
             connections.add(connection)

tests/rest/test_multiresult.py

@@ -1,19 +1,17 @@
 import pytest
 
+from openeo import BatchJob
 from openeo.rest._testing import DummyBackend
 from openeo.rest.multiresult import MultiResult
 
 
 class TestMultiResultHandling:
-
-    def test_create_job_method(self, dummy_backend):
-        con = dummy_backend.connection
-        cube = con.load_collection("S2")
+    def test_flat_graph(self, dummy_backend):
+        cube = dummy_backend.connection.load_collection("S2")
         save1 = cube.save_result(format="GTiff")
         save2 = cube.save_result(format="netCDF")
         multi_result = MultiResult([save1, save2])
-        multi_result.create_job()
-        assert dummy_backend.get_batch_pg() == {
+        assert multi_result.flat_graph() == {
             "loadcollection1": {
                 "process_id": "load_collection",
                 "arguments": {"id": "S2", "spatial_extent": None, "temporal_extent": None},
@@ -29,7 +27,37 @@ def test_create_job_method(self, dummy_backend):
             },
         }
 
-    def test_create_job_on_connection(self, con120, dummy_backend):
+    def test_create_job_method(self, dummy_backend):
+        con = dummy_backend.connection
+        cube = con.load_collection("S2")
+        save1 = cube.save_result(format="GTiff")
+        save2 = cube.save_result(format="netCDF")
+        multi_result = MultiResult([save1, save2])
+        multi_result.create_job(title="multi result test")
+        assert dummy_backend.batch_jobs == {
+            "job-000": {
+                "job_id": "job-000",
+                "pg": {
+                    "loadcollection1": {
+                        "process_id": "load_collection",
+                        "arguments": {"id": "S2", "spatial_extent": None, "temporal_extent": None},
+                    },
+                    "saveresult1": {
+                        "process_id": "save_result",
+                        "arguments": {"data": {"from_node": "loadcollection1"}, "format": "GTiff", "options": {}},
+                    },
+                    "saveresult2": {
+                        "process_id": "save_result",
+                        "arguments": {"data": {"from_node": "loadcollection1"}, "format": "netCDF", "options": {}},
+                        "result": True,
+                    },
+                },
+                "status": "created",
+                "title": "multi result test",
+            }
+        }
+
+    def test_create_job_through_connection(self, con120, dummy_backend):
         con = dummy_backend.connection
         cube = con120.load_collection("S2")
         save1 = cube.save_result(format="GTiff")
@@ -51,3 +79,34 @@ def test_create_job_on_connection(self, con120, dummy_backend):
                 "result": True,
             },
         }
+
+    def test_execute_batch(self, dummy_backend):
+        con = dummy_backend.connection
+        cube = con.load_collection("S2")
+        save1 = cube.save_result(format="GTiff")
+        save2 = cube.save_result(format="netCDF")
+        multi_result = MultiResult([save1, save2])
+        job = multi_result.execute_batch(title="multi result test")
+        assert isinstance(job, BatchJob)
+        assert dummy_backend.batch_jobs == {
+            "job-000": {
+                "job_id": "job-000",
+                "pg": {
+                    "loadcollection1": {
+                        "process_id": "load_collection",
+                        "arguments": {"id": "S2", "spatial_extent": None, "temporal_extent": None},
+                    },
+                    "saveresult1": {
+                        "process_id": "save_result",
+                        "arguments": {"data": {"from_node": "loadcollection1"}, "format": "GTiff", "options": {}},
+                    },
+                    "saveresult2": {
+                        "process_id": "save_result",
+                        "arguments": {"data": {"from_node": "loadcollection1"}, "format": "netCDF", "options": {}},
+                        "result": True,
+                    },
+                },
+                "status": "finished",
+                "title": "multi result test",
+            }
+        }