test: fix coverage report test for bulk enhance

src/diffbot_kg/clients/base.py

@@ -41,6 +41,8 @@ def _merge_params(self, params) -> dict[str, Any]:
 
         params = params or {}
         params = {**self.default_params, **params}
+
+        # sourcery skip: inline-immediately-returned-variable
         params = {k: v for k, v in params.items() if v is not None}
         return params
 
@@ -59,9 +61,11 @@ async def _get(
             BaseDiffbotResponse: The response from the API.
         """
 
-        headers = {"accept": "application/json", **(headers or {})}
+        headers = headers or {}
 
         params = self._merge_params(params)
+
+        # sourcery skip: inline-immediately-returned-variable
         resp = await self.s.get(url, params=params, headers=headers)
         return resp
 
@@ -88,10 +92,10 @@ async def _post(
 
         headers = {
             "content-type": "application/json",
-            "accept": "application/json",
             **(headers or {}),
         }
 
+        # sourcery skip: inline-immediately-returned-variable
         resp = await self.s.post(url, params=params, headers=headers, json=json)
         return resp
 

src/diffbot_kg/clients/enhance.py

@@ -125,6 +125,7 @@ async def bulkjob_coverage_report(
         url = self.bulk_job_coverage_report_url.human_repr().format(
             bulkjobId=bulkjobId, reportId=reportId
         )
+
         resp = await self._get(url)
         resp.__class__ = DiffbotCoverageReportResponse
         return cast(DiffbotCoverageReportResponse, resp)

tests/functional/clients/test_enhance_client.py

@@ -3,8 +3,8 @@
 
 import pytest
 from aiohttp import ClientResponseError
-
 from diffbot_kg.clients.enhance import DiffbotEnhanceClient
+
 from tests.functional.conftest import (
     ORG2_ENTITY_ID,
     ORG2_NAME,
@@ -89,15 +89,21 @@ async def test_bulkjob_status(self, request, token: Secret):
 
         job_id = _get_job_id(request)
 
-        DELAY = 10
+        TIMEOUT = 60
+        BACKOFF_FACTOR = 1.5
+        backoff = 1
         start = time.time()
 
         # ACT
-        while time.time() - start <= DELAY:
+        while True:
             response = await client.bulkjob_status(job_id)
             if response.complete:
                 break
-            time.sleep(1)
+            elif time.time() - start > TIMEOUT:
+                pytest.fail("Bulk job status check did not complete in time")
+
+            time.sleep(backoff)
+            backoff *= BACKOFF_FACTOR
 
         # ASSERT
         assert response.status == 200
@@ -159,45 +165,54 @@ async def test_single_bulkjob_result(self, request, token: Secret):
         await client.close()
 
     @pytest.mark.asyncio
-    async def test_bulkjob_stop(self, request, token: Secret):
+    async def test_bulkjob_coverage_report(self, request, token: Secret):
         # ARRANGE
         client = DiffbotEnhanceClient(token=token.value)
 
         job_id = _get_job_id(request)
+        report_id = request.config.cache.get("enhanceBulkJobCoverageReportId", None)
+        if report_id is None:
+            pytest.fail("Enhance bulk job coverage report ID not found in cache")
+
+
+        TIMEOUT = 60
+        BACKOFF_FACTOR = 1.5
+        backoff = 1
+        start = time.time()
 
         # ACT
-        response = await client.stop_bulkjob(job_id)
+        while True:
+            try:
+                response = await client.bulkjob_coverage_report(job_id, report_id)
+            except ClientResponseError as e:
+                if e.status == 400:
+                    time.sleep(backoff)
+                    backoff *= BACKOFF_FACTOR
+            else:
+                if response.status == 200:
+                    break
+                elif time.time() - start > TIMEOUT:
+                    pytest.fail("Bulk job coverage report did not generate in time")
 
         # ASSERT
         assert response.status == 200
-        assert response.content["status"] == "COMPLETE"
-        assert response.content["message"] == f"Bulkjob [{job_id}] is completed"
+        assert len(response.content.strip().split("\n")) == 4
 
-        # TEARDOWN
-        await client.close()
 
     @pytest.mark.asyncio
-    async def test_bulkjob_coverage_report(self, request, token: Secret):
+    async def test_bulkjob_stop(self, request, token: Secret):
         # ARRANGE
         client = DiffbotEnhanceClient(token=token.value)
 
         job_id = _get_job_id(request)
-        report_id = request.config.cache.get("enhanceBulkJobCoverageReportId", None)
-        if report_id is None:
-            pytest.fail("Enhance bulk job coverage report ID not found in cache")
-
-        DELAY = 10
-        start = time.time()
 
         # ACT
-        while time.time() - start <= DELAY:
-            try:
-                response = await client.bulkjob_coverage_report(job_id, report_id)
-            except ClientResponseError:
-                time.sleep(1)
-            else:
-                break
+        response = await client.stop_bulkjob(job_id)
 
         # ASSERT
         assert response.status == 200
-        assert len(response.content.strip().split("\n")) == 4
+        assert response.content["status"] == "COMPLETE"
+        assert response.content["message"] == f"Bulkjob [{job_id}] is completed"
+
+        # TEARDOWN
+        await client.close()

tests/unit/clients/test_enhance_client.py

@@ -33,7 +33,7 @@ async def test_mocked_enhance(self, mocker, client):
         DiffbotSession.get.assert_called_with(
             DiffbotEnhanceClient.enhance_url,
             params={**params, "token": "valid_token"},
-            headers={"accept": "application/json"},
+            headers={}
         )
         assert isinstance(response, DiffbotEntitiesResponse)
         assert response.status == 200
@@ -59,7 +59,7 @@ async def test_mocked_create_bulkjob(self, mocker, client):
         DiffbotSession.post.assert_called_with(
             DiffbotEnhanceClient.bulk_job_url,
             params={"token": "valid_token"},
-            headers={"content-type": "application/json", "accept": "application/json"},
+            headers={"content-type": "application/json"},
             json=params
         )
         assert isinstance(response, DiffbotBulkJobCreateResponse)

tests/unit/clients/test_search_client.py

@@ -37,7 +37,7 @@ async def test_mocked_search(self, mocker, client):
         DiffbotSession.get.assert_called_with(
             DiffbotSearchClient.search_url,
             params=params,
-            headers={"accept": "application/json"},
+            headers={},
         )
         assert isinstance(response, DiffbotEntitiesResponse)
         assert response.status == 200