PT5353 Adding collection results to processing Job (#632)

* adding collection to processing Job

* updating CHANGELOG

* fix pylint collection issue

* adding collection as a property

* adding StacClient descriptor to base module

* adding job.collection tests

* adding job.collection with None tests

* adding coverage for StacClient descriptor

* updating changelog

* updating stac_client in Job

* second round of feedback

* Address comments

---------

Co-authored-by: Javid Gafar-zada <[email protected]>

CHANGELOG.md

@@ -29,6 +29,13 @@ You can check your current version with the following command:
     ```
 
 For more information, see [UP42 Python package description](https://pypi.org/project/up42-py/).
+## 1.0.4a17
+
+**Jun 13, 2024**
+
+- Added `collection` to job class in processing module.
+- Added `StacClient` descriptor to base module.
+
 ## 1.0.4a16
 
 **Jun 11, 2024**

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "up42-py"
-version = "1.0.4a16"
+version = "1.0.4a17"
 description = "Python SDK for UP42, the geospatial marketplace and developer platform."
 authors = ["UP42 GmbH <[email protected]>"]
 license = "https://github.com/up42/up42-py/blob/master/LICENSE"

tests/test_base.py

@@ -2,7 +2,9 @@
 from typing import Union
 from unittest import mock
 
+import pystac_client
 import pytest
+import requests_mock as req_mock
 
 from up42 import base
 
@@ -40,6 +42,7 @@ class ActiveRecord:
     session = base.Session()
     class_workspace_id = base.WorkspaceId()
     workspace_id: Union[str, base.WorkspaceId] = dataclasses.field(default=base.WorkspaceId())
+    stac_client = base.StacClient()
 
 
 class TestDescriptors:
@@ -69,3 +72,9 @@ def test_workspace_id_should_provide_default_value(self):
         assert record.workspace_id == constants.WORKSPACE_ID
         assert ActiveRecord.class_workspace_id == constants.WORKSPACE_ID
         assert constants.WORKSPACE_ID in repr(record)
+
+    def test_should_provide_stac_client(self, requests_mock: req_mock.Mocker):
+        requests_mock.get(constants.URL_STAC_CATALOG, json=constants.STAC_CATALOG_RESPONSE)
+        record = ActiveRecord()
+        assert isinstance(record.stac_client, pystac_client.Client)
+        assert requests_mock.called

tests/test_processing.py

@@ -20,7 +20,9 @@
 COST_URL = f"{constants.API_HOST}/v2/processing/processes/{PROCESS_ID}/cost"
 EXECUTION_URL = f"{constants.API_HOST}/v2/processing/processes/{PROCESS_ID}/execution"
 TITLE = "title"
-ITEM_URL = "https://item-url"
+COLLECTION_ID = str(uuid.uuid4())
+COLLECTION_URL = f"https://collections/{COLLECTION_ID}"
+ITEM_URL = "https://item-url/"
 ITEM = pystac.Item.from_dict(
     {
         "type": "Feature",
@@ -53,6 +55,7 @@
     "accountID": ACCOUNT_ID,
     "workspaceID": constants.WORKSPACE_ID,
     "definition": DEFINITION,
+    "results": {"collection": f"{COLLECTION_URL}"},
     "status": "created",
     "created": f"{NOW.isoformat()}Z",
     "updated": f"{NOW.isoformat()}Z",
@@ -66,19 +69,19 @@
     account_id=ACCOUNT_ID,
     workspace_id=constants.WORKSPACE_ID,
     definition=DEFINITION,
+    collection_url=COLLECTION_URL,
     status=processing.JobStatus.CREATED,
     created=NOW,
     updated=NOW,
 )
 
 
 @pytest.fixture(autouse=True)
-def workspace():
-    with mock.patch("up42.base.workspace") as workspace_mock:
-        session = requests.Session()
-        session.hooks = {"response": lambda response, *args, **kwargs: response.raise_for_status()}
-        workspace_mock.auth.session = session
-        yield
+def set_status_raising_session():
+    session = requests.Session()
+    session.hooks = {"response": lambda response, *args, **kwargs: response.raise_for_status()}
+    processing.Job.session = session  # type: ignore
+    processing.JobTemplate.session = session  # type: ignore
 
 
 class TestCost:
@@ -286,6 +289,15 @@ def test_should_get_job(self, requests_mock: req_mock.Mocker):
         requests_mock.get(url=JOB_URL, json=JOB_METADATA)
         assert processing.Job.get(JOB_ID) == JOB
 
+    def test_should_get_collection(self):
+        stac_client = JOB.stac_client = mock.MagicMock()
+        assert JOB.collection == stac_client.get_collection.return_value
+        stac_client.get_collection.assert_called_with(COLLECTION_ID)
+
+    def test_should_get_no_collection_if_collection_url_is_missing(self):
+        job = dataclasses.replace(JOB, collection_url=None)
+        assert not job.collection
+
     @pytest.mark.parametrize("process_id", [None, [PROCESS_ID]])
     @pytest.mark.parametrize("workspace_id", [None, constants.WORKSPACE_ID])
     @pytest.mark.parametrize("status", [None, [processing.JobStatus.CAPTURED]])

up42/base.py

@@ -3,6 +3,7 @@
 import warnings
 from typing import Any, Optional, Union
 
+import pystac_client
 import requests
 
 from up42 import auth as up42_auth
@@ -92,3 +93,8 @@ def __set__(self, obj, value: str) -> None:
         if value == self:
             value = workspace.id
         obj.__dict__["workspace_id"] = value
+
+
+class StacClient:
+    def __get__(self, obj, obj_type=None) -> pystac_client.Client:
+        return utils.stac_client(workspace.auth.client.auth)

up42/processing.py

@@ -29,12 +29,17 @@ class JobStatus(enum.Enum):
     RELEASED = "released"
 
 
+class JobResults(TypedDict, total=False):
+    collection: Optional[str]
+
+
 class JobMetadata(TypedDict):
     processID: str  # pylint: disable=invalid-name
     jobID: str  # pylint: disable=invalid-name
     accountID: str  # pylint: disable=invalid-name
     workspaceID: Optional[str]  # pylint: disable=invalid-name
     definition: dict
+    results: JobResults
     status: str
     created: str
     started: Optional[str]
@@ -56,6 +61,7 @@ def _to_datetime(value: Optional[str]):
 @dataclasses.dataclass
 class Job:
     session = base.Session()
+    stac_client = base.StacClient()
     process_id: str
     id: str
     account_id: str
@@ -64,16 +70,25 @@ class Job:
     status: JobStatus
     created: datetime.datetime
     updated: datetime.datetime
+    collection_url: Optional[str] = None
     started: Optional[datetime.datetime] = None
     finished: Optional[datetime.datetime] = None
 
+    @property
+    def collection(self) -> Optional[pystac.Collection]:
+        if self.collection_url is None:
+            return None
+        collection_id = self.collection_url.split("/")[-1]
+        return self.stac_client.get_collection(collection_id)
+
     @staticmethod
     def from_metadata(metadata: JobMetadata) -> "Job":
         return Job(
             process_id=metadata["processID"],
             id=metadata["jobID"],
             account_id=metadata["accountID"],
             workspace_id=metadata["workspaceID"],
+            collection_url=metadata.get("results", {}).get("collection"),
             definition=metadata["definition"],
             status=JobStatus(metadata["status"]),
             created=_to_datetime(metadata["created"]),