RestApiConnection.request: add option for expected status code

Author: soxofaan

openeo/rest/connection.py

@@ -12,6 +12,7 @@
 import requests
 from deprecated import deprecated
 from openeo.rest import OpenEoClientException
+from openeo.util import ensure_list
 from requests import Response
 from requests.auth import HTTPBasicAuth, AuthBase
 
@@ -67,7 +68,8 @@ def _merged_headers(self, headers: dict) -> dict:
             result.update(headers)
         return result
 
-    def request(self, method: str, path: str, headers: dict = None, auth: AuthBase = None, check_status=True, **kwargs):
+    def request(self, method: str, path: str, headers: dict = None, auth: AuthBase = None,
+                check_error=True, expected_status=None, **kwargs):
         """Generic request send"""
         resp = self.session.request(
             method=method,
@@ -76,10 +78,12 @@ def request(self, method: str, path: str, headers: dict = None, auth: AuthBase =
             auth=auth or self.auth,
             **kwargs
         )
-        if check_status:
-            # TODO: option to specify the list/range of expected status codes?
-            if resp.status_code >= 400:
-                self._raise_api_error(resp)
+        # Check for API errors and unexpected HTTP status codes as desired.
+        status = resp.status_code
+        if check_error and status >= 400:
+            self._raise_api_error(resp)
+        if expected_status and status not in ensure_list(expected_status):
+            raise OpenEoClientException("Status code {s} is not expected {e}".format(s=status, e=expected_status))
         return resp
 
     def _raise_api_error(self, response: requests.Response):

openeo/util.py

@@ -34,6 +34,14 @@ def first_not_none(*args):
     raise ValueError("No not-None values given.")
 
 
+def ensure_list(x):
+    """Convert given data structure to a list."""
+    try:
+        return list(x)
+    except TypeError:
+        return [x]
+
+
 def get_temporal_extent(*args,
                         start_date: Union[str, datetime, date] = None, end_date: Union[str, datetime, date] = None,
                         extent: Union[list, tuple] = None,

tests/rest/test_connection.py

@@ -2,6 +2,7 @@
 
 import pytest
 import requests_mock
+from openeo.rest import OpenEoClientException
 
 from openeo.rest.auth.auth import NullAuth, BearerAuth
 from openeo.rest.connection import Connection, RestApiConnection, connect, OpenEoApiError
@@ -48,6 +49,35 @@ def text(request, context):
         conn.post("/foo", {}, headers={"X-Openeo-Bar": "XY123"})
 
 
+def test_rest_api_expected_status(requests_mock):
+    conn = RestApiConnection(API_URL)
+    requests_mock.get("https://oeo.net/foo", status_code=200, json={"o": "k"})
+    # Expected status
+    assert conn.get("/foo", expected_status=200).json() == {"o": "k"}
+    assert conn.get("/foo", expected_status=[200, 201]).json() == {"o": "k"}
+    # Unexpected status
+    with pytest.raises(OpenEoClientException, match="Status code 200 is not expected 204"):
+        conn.get("/foo", expected_status=204)
+    with pytest.raises(OpenEoClientException, match=r"Status code 200 is not expected \[203, 204\]"):
+        conn.get("/foo", expected_status=[203, 204])
+
+
+def test_rest_api_expected_status_with_error(requests_mock):
+    conn = RestApiConnection(API_URL)
+    requests_mock.get("https://oeo.net/bar", status_code=406, json={"code": "NoBar", "message": "no bar please"})
+    # First check for API error by default
+    with pytest.raises(OpenEoApiError, match=r"\[406\] NoBar: no bar please"):
+        conn.get("/bar", expected_status=200)
+    with pytest.raises(OpenEoApiError, match=r"\[406\] NoBar: no bar please"):
+        conn.get("/bar", expected_status=[201, 202])
+    # Don't check for error, just status
+    conn.get("/bar", check_error=False, expected_status=406)
+    with pytest.raises(OpenEoClientException, match="Status code 406 is not expected 302"):
+        conn.get("/bar", check_error=False, expected_status=302)
+    with pytest.raises(OpenEoClientException, match=r"Status code 406 is not expected \[302, 303\]"):
+        conn.get("/bar", check_error=False, expected_status=[302, 303])
+
+
 def test_connection_with_session():
     session = mock.Mock()
     response = session.request.return_value

tests/test_util.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from openeo.util import first_not_none, get_temporal_extent, TimingLogger
+from openeo.util import first_not_none, get_temporal_extent, TimingLogger, ensure_list
 
 
 @pytest.mark.parametrize(['input', 'expected'], [
@@ -30,6 +30,19 @@ def test_first_not_none_failures():
         first_not_none(None, None)
 
 
+@pytest.mark.parametrize(["input", "expected"], [
+    (None, [None]),
+    (123, [123]),
+    ("abc", ["a", "b", "c"]),
+    ([1, 2, "three"], [1, 2, "three"]),
+    ((1, 2, "three"), [1, 2, "three"]),
+    ({1: "a", 2: "b"}, [1, 2]),
+    ({1, 2, 3, 3, 2}, [1, 2, 3]),
+])
+def test_ensure_list(input, expected):
+    assert ensure_list(input) == expected
+
+
 def test_get_temporal_extent():
     assert get_temporal_extent("2019-03-15") == ("2019-03-15", None)
     assert get_temporal_extent("2019-03-15", "2019-10-11") == ("2019-03-15", "2019-10-11")