From c24515054a7ee3cf95590f409af0c0d3ad30ce46 Mon Sep 17 00:00:00 2001 From: Brendan Smith Date: Fri, 26 Apr 2024 00:06:47 -0500 Subject: [PATCH] test: fix coverage report test for bulk enhance --- src/diffbot_kg/clients/base.py | 8 ++- src/diffbot_kg/clients/enhance.py | 1 + .../functional/clients/test_enhance_client.py | 65 ++++++++++++------- tests/unit/clients/test_enhance_client.py | 4 +- tests/unit/clients/test_search_client.py | 2 +- 5 files changed, 50 insertions(+), 30 deletions(-) diff --git a/src/diffbot_kg/clients/base.py b/src/diffbot_kg/clients/base.py index 99292fa..85ad949 100644 --- a/src/diffbot_kg/clients/base.py +++ b/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 diff --git a/src/diffbot_kg/clients/enhance.py b/src/diffbot_kg/clients/enhance.py index e103d46..1debead 100644 --- a/src/diffbot_kg/clients/enhance.py +++ b/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) diff --git a/tests/functional/clients/test_enhance_client.py b/tests/functional/clients/test_enhance_client.py index ee04a2b..0086a53 100644 --- a/tests/functional/clients/test_enhance_client.py +++ b/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() diff --git a/tests/unit/clients/test_enhance_client.py b/tests/unit/clients/test_enhance_client.py index 40c13b6..1db34c4 100644 --- a/tests/unit/clients/test_enhance_client.py +++ b/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) diff --git a/tests/unit/clients/test_search_client.py b/tests/unit/clients/test_search_client.py index 7941d42..aa055ae 100644 --- a/tests/unit/clients/test_search_client.py +++ b/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