more tests for STACAPIJobDatabase

Author: VincentVerelst

openeo/extra/stac_job_db.py

@@ -216,7 +216,7 @@ def _ingest_bulk(self, items: Iterable[Item]) -> dict:
 
         url_path = f"collections/{collection_id}/bulk_items"
         data = {"method": "upsert", "items": {item.id: item.to_dict() for item in items}}
-        response = requests.post(self.join_url(url_path), auth=self._auth, json=data)
+        response = requests.post(url=self.join_url(url_path), auth=self._auth, json=data)
 
         _log.info(f"HTTP response: {response.status_code} - {response.reason}: body: {response.json()}")
 
@@ -230,7 +230,6 @@ def _upload_items_bulk(self, collection_id: str, items: Iterable[Item]) -> None:
         with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
             for item in items:
                 self._prepare_item(item, collection_id)
-                # item.validate()
                 chunk.append(item)
 
                 if len(chunk) == self.bulk_size:

tests/extra/test_stac_jobdb.py

@@ -25,14 +25,18 @@ def mock_stac_api_job_database(mock_auth) -> STACAPIJobDatabase:
 
 
 @pytest.fixture
-def mock_pystac_client():
+def mock_pystac_client(dummy_stac_item):
     mock_client = MagicMock(spec=pystac_client.Client)
 
     mock_client.get_collections.return_value = [
         MagicMock(id="collection-1"),
         MagicMock(id="collection-2"),
     ]
 
+    mock_item_search = MagicMock(spec=pystac_client.ItemSearch)
+    mock_item_search.items.return_value = [dummy_stac_item]
+    mock_client.search.return_value = mock_item_search
+
     with patch("pystac_client.Client.open", return_value=mock_client):
         yield mock_client
 
@@ -195,6 +199,16 @@ def patch_datetime_now():
         yield mock_datetime
 
 
+@pytest.fixture
+def bulk_dataframe():
+    return pd.DataFrame(
+        {
+            "some_property": [f"value-{i}" for i in range(10)],
+        },
+        index=[i for i in range(10)],
+    )
+
+
 class TestSTACAPIJobDatabase:
     def test_exists(self, job_db_exists, job_db_not_exists):
 
@@ -246,7 +260,6 @@ def test_initialize_from_df_with_geometry(
         assert job_db_not_exists.geometry_column == "geometry"
 
     def test_series_from(self, job_db_exists, dummy_series, dummy_stac_item):
-        job_db_exists.has_geometry = True
         pdt.assert_series_equal(job_db_exists.series_from(dummy_stac_item), dummy_series)
 
     def test_item_from(self, patch_datetime_now, job_db_exists, dummy_series, dummy_stac_item):
@@ -260,12 +273,119 @@ def test_item_from_geometry(
         item = job_db_exists.item_from(dummy_series_geometry)
         assert item.to_dict() == dummy_stac_item_geometry.to_dict()
 
-    def test_count_by_status(self):
-        pass
-
-    def test_get_by_status(self):
-        pass
+    @patch("openeo.extra.stac_job_db.STACAPIJobDatabase.get_by_status")
+    def test_count_by_status(self, mock_get_by_status, normalized_dummy_dataframe, job_db_exists):
+        mock_get_by_status.return_value = normalized_dummy_dataframe
+        assert job_db_exists.count_by_status() == {"not_started": 1}
+
+    def test_get_by_status_no_filter(self, job_db_exists):
+        job_db_exists.get_by_status(())
+
+        job_db_exists.client.search.assert_called_once_with(
+            method="GET", collections=["collection-1"], filter=None, max_items=None
+        )
+
+    def test_get_by_status_with_filter(self, job_db_exists):
+        job_db_exists.get_by_status(["not_started"])
+        job_db_exists.client.search.assert_called_once_with(
+            method="GET", collections=["collection-1"], filter="\"properties.status\"='not_started'", max_items=None
+        )
+
+    def test_get_by_status_result(self, job_db_exists):
+        df = job_db_exists.get_by_status(["not_started"])
+
+        pdt.assert_frame_equal(
+            df,
+            pd.DataFrame(
+                {
+                    "datetime": [pystac.utils.datetime_to_str(FAKE_NOW)],
+                    "some_property": ["value"],
+                },
+                index=["test"],
+            ),
+        )
+
+    @patch("requests.post")
+    def test_persist_single_chunk(self, mock_requests_post, bulk_dataframe, job_db_exists, patch_datetime_now):
+        def bulk_items(df):
+            all_items = []
+            if not df.empty:
+
+                def handle_row(series):
+                    item = job_db_exists.item_from(series)
+                    job_db_exists._prepare_item(item, job_db_exists.collection_id)
+                    all_items.append(item)
+
+                df.apply(handle_row, axis=1)
+            return all_items
+
+        items = bulk_items(bulk_dataframe)
+
+        mock_requests_post.return_value.status_code = 200
+        mock_requests_post.return_value.json.return_value = {"status": "success"}
+        mock_requests_post.reason = "OK"
+
+        job_db_exists.persist(bulk_dataframe)
+        # job_db_exists._upload_items_bulk(collection_id=job_db_exists.collection_id, items=items)
+
+        mock_requests_post.assert_called_once()
+
+        mock_requests_post.assert_called_with(
+            url=f"http://fake-stac-api/collections/{job_db_exists.collection_id}/bulk_items",
+            auth=None,
+            json={
+                "method": "upsert",
+                "items": {item.id: item.to_dict() for item in items},
+            },
+        )
+
+    @patch("requests.post")
+    def test_persist_multiple_chunks(self, mock_requests_post, bulk_dataframe, job_db_exists):
+        def bulk_items(df):
+            all_items = []
+            if not df.empty:
+
+                def handle_row(series):
+                    item = job_db_exists.item_from(series)
+                    job_db_exists._prepare_item(item, job_db_exists.collection_id)
+                    all_items.append(item)
+
+                df.apply(handle_row, axis=1)
+            return all_items
+
+        items = bulk_items(bulk_dataframe)
+
+        mock_requests_post.return_value.status_code = 200
+        mock_requests_post.return_value.json.return_value = {"status": "success"}
+        mock_requests_post.reason = "OK"
+
+        job_db_exists.bulk_size = 3
+        job_db_exists._upload_items_bulk(collection_id=job_db_exists.collection_id, items=items)
+
+        # 10 items in total, 3 items per chunk, should result in 4 calls
+        assert mock_requests_post.call_count == 4
+        expected_calls = [
+            {
+                "url": f"http://fake-stac-api/collections/{job_db_exists.collection_id}/bulk_items",
+                "auth": None,
+                "json": {"method": "upsert", "items": {item.id: item.to_dict() for item in items[:3]}},
+            },
+            {
+                "url": f"http://fake-stac-api/collections/{job_db_exists.collection_id}/bulk_items",
+                "auth": None,
+                "json": {"method": "upsert", "items": {item.id: item.to_dict() for item in items[3:6]}},
+            },
+            {
+                "url": f"http://fake-stac-api/collections/{job_db_exists.collection_id}/bulk_items",
+                "auth": None,
+                "json": {"method": "upsert", "items": {item.id: item.to_dict() for item in items[6:9]}},
+            },
+            {
+                "url": f"http://fake-stac-api/collections/{job_db_exists.collection_id}/bulk_items",
+                "auth": None,
+                "json": {"method": "upsert", "items": {item.id: item.to_dict() for item in items[9:]}},
+            },
+        ]
 
-    def test_persist(self):
-        pass
-        # This should test upload items bulk
+        for i, call in enumerate(mock_requests_post.call_args_list):
+            assert call[1] == expected_calls[i]