From 13e52b7e7e9706f295ead7c985124ee84ab6ec51 Mon Sep 17 00:00:00 2001 From: Robin5605 Date: Sat, 6 Apr 2024 01:43:43 -0500 Subject: [PATCH 1/2] Use httpx client in report endpoint Also make a dependency and use FastAPI dependency injection to accept the httpx.Client instance. This also makes it much easier to test. --- src/mainframe/dependencies.py | 6 ++++++ src/mainframe/endpoints/report.py | 14 +++++++++++--- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/src/mainframe/dependencies.py b/src/mainframe/dependencies.py index 32eb5ac0..99a2e461 100644 --- a/src/mainframe/dependencies.py +++ b/src/mainframe/dependencies.py @@ -12,6 +12,12 @@ from mainframe.rules import Rules +@cache +def get_http_client() -> httpx.Client: + http_client = httpx.Client() + return http_client + + @cache def get_pypi_client() -> PyPIServices: http_client = httpx.Client() diff --git a/src/mainframe/endpoints/report.py b/src/mainframe/endpoints/report.py index a5325884..8b383913 100644 --- a/src/mainframe/endpoints/report.py +++ b/src/mainframe/endpoints/report.py @@ -11,7 +11,7 @@ from mainframe.constants import mainframe_settings from mainframe.database import get_db -from mainframe.dependencies import get_pypi_client, validate_token +from mainframe.dependencies import get_http_client, get_pypi_client, validate_token from mainframe.json_web_token import AuthenticationData from mainframe.models.orm import Scan from mainframe.models.schemas import ( @@ -159,6 +159,7 @@ def report_package( session: Annotated[Session, Depends(get_db)], auth: Annotated[AuthenticationData, Depends(validate_token)], pypi_client: Annotated[PyPIServices, Depends(get_pypi_client)], + http_client: Annotated[httpx.Client, Depends(get_http_client)], ): """ Report a package to PyPI. @@ -218,7 +219,14 @@ def report_package( additional_information=body.additional_information, ) - httpx.post(f"{mainframe_settings.reporter_url}/report/email", json=jsonable_encoder(report)) + response = http_client.post(f"{mainframe_settings.reporter_url}/report/email", json=jsonable_encoder(report)) + try: + response.raise_for_status() + except httpx.HTTPStatusError as err: + detail = "Dragonfly Reporter service failed" + log.error(detail, status=err.response.status_code, message=err.response.text) + raise HTTPException(502, detail=detail) + else: # We previously checked this condition, but the typechecker isn't smart # enough to figure that out @@ -231,7 +239,7 @@ def report_package( extra=dict(yara_rules=rules_matched), ) - httpx.post(f"{mainframe_settings.reporter_url}/report/{name}", json=jsonable_encoder(report)) + http_client.post(f"{mainframe_settings.reporter_url}/report/{name}", json=jsonable_encoder(report)) log.info( "Sent report", From 805fdda6faa4eb5f54e524d8068d6e54bab1340d Mon Sep 17 00:00:00 2001 From: Robin5605 Date: Sat, 6 Apr 2024 01:44:43 -0500 Subject: [PATCH 2/2] Add test case for Reporter service failing --- tests/test_report.py | 49 +++++++++++++++++++++++++++++++++++++++----- 1 file changed, 44 insertions(+), 5 deletions(-) diff --git a/tests/test_report.py b/tests/test_report.py index 78365757..c72ed77a 100644 --- a/tests/test_report.py +++ b/tests/test_report.py @@ -1,6 +1,6 @@ -from datetime import datetime, timedelta +from datetime import UTC, datetime, timedelta from typing import Optional -from unittest.mock import MagicMock +from unittest.mock import Mock import httpx import pytest @@ -105,11 +105,12 @@ def test_report( db_session.add(scan) db_session.commit() - httpx.post = MagicMock() + mock_http_client = Mock(spec=httpx.Client) + mock_http_client.post = Mock() - report_package(body, db_session, auth, pypi_client) + report_package(body, db_session, auth, pypi_client, mock_http_client) - httpx.post.assert_called_once_with(url, json=jsonable_encoder(expected)) + mock_http_client.post.assert_called_once_with(url, json=jsonable_encoder(expected)) scan = db_session.scalar(select(Scan).where(Scan.name == "c").where(Scan.version == "1.0.0")) @@ -412,3 +413,41 @@ def test_report_lookup_package(db_session: Session): res = _lookup_package("c", "1.0.0", db_session) assert res == scan + + +def test_reporter_service_fail(db_session: Session, auth: AuthenticationData, pypi_client: PyPIServices): + scan = Scan( + name="abc", + version="1.0.0", + status=Status.FINISHED, + inspector_url="test inspector URL", + score=25, + queued_at=datetime.now(UTC) - timedelta(seconds=30), + queued_by="remmy", + ) + db_session.add(scan) + db_session.commit() + + mock_http_client = Mock(spec=httpx.Client) + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 400 + + def side_effect(): + raise httpx.HTTPStatusError("Error", request=Mock(), response=mock_response) + + mock_response.raise_for_status = Mock(side_effect=side_effect) + mock_http_client.post = Mock(return_value=mock_response) + + body = ReportPackageBody( + name="abc", + version="1.0.0", + recipient=None, + inspector_url="test inspector url", + additional_information="this is a bad package", + use_email=True, + ) + + with pytest.raises(HTTPException) as err: + report_package(body, db_session, auth, pypi_client, mock_http_client) + + assert err.value.status_code == 502