RESTJob.download_result and execute_batch: finetune and unit testing

Author: soxofaan

openeo/rest/imagecollectionclient.py

@@ -969,7 +969,7 @@ def execute_batch(
             self,
             outputfile: str, out_format: str = None,
             print=print, max_poll_interval=60, connection_retry_interval=30,
-            job_options = None,**format_options):
+            job_options=None, **format_options):
         """
         Evaluate the process graph by creating a batch job, and retrieving the results when it is finished.
         This method is mostly recommended if the batch job is expected to run in a reasonable amount of time.
@@ -986,12 +986,13 @@ def execute_batch(
         job = self.send_job(out_format,job_options=job_options, **format_options)
         job.start_job()
 
-        job_id = None
+        job_id = job.job_id
         job_info = None
         status = None
-        poll_interval = 5
+        poll_interval = min(5, max_poll_interval)
         start_time = time.time()
         while True:
+            # TODO: also allow a hard time limit on this infinite poll loop?
             elapsed = str(datetime.timedelta(seconds=time.time() - start_time))
             try:
                 job_info = job.describe_job()
@@ -1000,11 +1001,11 @@ def execute_batch(
                 time.sleep(connection_retry_interval)
                 continue
 
-            job_id = job_info.get("id", "N/A")
             status = job_info.get("status", "N/A")
             print("{t} Job {i}: {s} (progress {p})".format(
-                t=elapsed, i=job_id, s=status, p=job_info.get("progress", "N/A"))
-            )
+                t=elapsed, i=job_id, s=status,
+                p='{p}%'.format(p=job_info["progress"]) if "progress" in job_info else "N/A"
+            ))
             if status not in ('submitted', 'queued', 'running'):
                 break
 

openeo/rest/job.py

@@ -1,14 +1,18 @@
+import logging
 import typing
 import urllib.request
 from typing import List
 
 from openeo.job import Job, JobResult
+from openeo.rest import OpenEoClientException
 
 if hasattr(typing, 'TYPE_CHECKING') and typing.TYPE_CHECKING:
     # Only import this for type hinting purposes. Runtime import causes circular dependency issues.
     # Note: the `hasattr` check is necessary for Python versions before 3.5.2.
     from openeo.rest.connection import Connection
 
+logger = logging.getLogger(__name__)
+
 
 class RESTJobResult(JobResult):
     def __init__(self, url):
@@ -54,9 +58,10 @@ def estimate_job(self):
     def start_job(self):
         """ Start / queue a job for processing."""
         # POST /jobs/{job_id}/results
-        request = self.connection.post("/jobs/{}/results".format(self.job_id))
-
-        return request.status_code
+        url = "/jobs/{}/results".format(self.job_id)
+        request = self.connection.post(url)
+        if request.status_code != 202:
+            logger.warning("{u} returned with status code {s} instead of 202".format(u=url, s=request.status_code))
 
     def stop_job(self):
         """ Stop / cancel job processing."""
@@ -72,34 +77,24 @@ def list_results(self, type=None):
 
     def download_results(self, target):
         """ Download job results."""
-        # GET /jobs/{job_id}/results > ...
-
-        download_url = "/jobs/{}/results".format(self.job_id)
-        r = self.connection.get(download_url, stream=True)
-
-        if r.status_code == 200:
-
-            url = r.json()
-            if "links" in url:
-                download_url = url["links"][0]
-                if "href" in download_url:
-                    download_url = download_url["href"]
-
-            with open(target, 'wb') as handle:
-                response = self.connection.get(download_url, stream=True)
-
-                for block in response.iter_content(1024):
-
-                    if not block:
-                        break
-
-                    handle.write(block)
-        else:
-            raise ConnectionAbortedError(r.text)
-        return r.status_code
-
-# TODO: All below methods are deprecated (at least not specified in the coreAPI)
-    def download(self, outputfile:str, outputformat=None):
+        results_url = "/jobs/{}/results".format(self.job_id)
+        r = self.connection.get(results_url, expected_status=200)
+        links = r.json()["links"]
+        if len(links) != 1:
+            # TODO handle download of multiple files?
+            raise OpenEoClientException("Expected 1 result file to download, but got {c}".format(c=len(links)))
+        file_url = links[0]["href"]
+
+        with open(target, 'wb') as handle:
+            response = self.connection.get(file_url, stream=True)
+            for block in response.iter_content(1024):
+                if not block:
+                    break
+                handle.write(block)
+        return target
+
+    # TODO: All below methods are deprecated (at least not specified in the coreAPI)
+    def download(self, outputfile: str, outputformat=None):
         """ Download the result as a raster."""
         try:
             return self.connection.download_job(self.job_id, outputfile, outputformat)
@@ -117,6 +112,3 @@ def queue(self):
     def results(self) -> List[RESTJobResult]:
         """ Returns this job's results. """
         return [RESTJobResult(link['href']) for link in self.connection.job_results(self.job_id)['links']]
-
-
-

tests/rest/test_job.py

@@ -0,0 +1,49 @@
+import re
+
+import openeo
+import pytest
+
+API_URL = "https://oeo.net"
+
+
+@pytest.fixture
+def session040(requests_mock):
+    requests_mock.get(API_URL + "/", json={"api_version": "0.4.0"})
+    session = openeo.connect(API_URL)
+    return session
+
+
+def test_execute_batch(session040, requests_mock, tmpdir):
+    requests_mock.get(API_URL + "/collections/SENTINEL2", json={"foo": "bar"})
+    requests_mock.post(API_URL + "/jobs", headers={"OpenEO-Identifier": "f00ba5"})
+    requests_mock.post(API_URL + "/jobs/f00ba5/results")
+    requests_mock.get(API_URL + "/jobs/f00ba5", [
+        {'json': {"status": "submitted"}},
+        {'json': {"status": "queued"}},
+        {'json': {"status": "running", "progress": 15}},
+        {'json': {"status": "running", "progress": 80}},
+        {'json': {"status": "finished", "progress": 100}},
+    ])
+    requests_mock.get(API_URL + "/jobs/f00ba5/results", json={
+        "links": [{"href": API_URL + "/jobs/f00ba5/files/output.tiff"}]
+    })
+    requests_mock.get(API_URL + "/jobs/f00ba5/files/output.tiff", text="tiffdata")
+
+    path = tmpdir.join("tmp.tiff")
+    log = []
+
+    def print(msg):
+        log.append(msg)
+
+    session040.load_collection("SENTINEL2").execute_batch(
+        outputfile=str(path), out_format="GTIFF",
+        max_poll_interval=.1, print=print
+    )
+
+    assert re.match(r"0:00:00(.0\d*)? Job f00ba5: submitted \(progress N/A\)", log[0])
+    assert re.match(r"0:00:00.1\d* Job f00ba5: queued \(progress N/A\)", log[1])
+    assert re.match(r"0:00:00.2\d* Job f00ba5: running \(progress 15%\)", log[2])
+    assert re.match(r"0:00:00.3\d* Job f00ba5: running \(progress 80%\)", log[3])
+    assert re.match(r"0:00:00.4\d* Job f00ba5: finished \(progress 100%\)", log[4])
+
+    assert path.read() == "tiffdata"