Mypy fix

ThePalaceProject · Oct 21, 2024 · e826ab3 · e826ab3
1 parent 1630416
commit e826ab3
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 13 deletions.
diff --git a/src/palace/manager/api/controller/odl_notification.py b/src/palace/manager/api/controller/odl_notification.py
@@ -5,7 +5,7 @@
 from flask_babel import lazy_gettext as _
 from pydantic import ValidationError
 from sqlalchemy import select
-from sqlalchemy.orm import Session, object_session
+from sqlalchemy.orm import Session
 from sqlalchemy.orm.exc import StaleDataError
 
 from palace.manager.api.odl.api import OPDS2WithODLApi
@@ -102,9 +102,8 @@ def _process_notification(self, loan: Loan | None) -> Response:
             #   Once we move the OPDS2WithODL scripts to celery this should be possible.
             #   For now we just mark the loan as expired.
             if not status_doc.active:
-                session = object_session(loan)
                 try:
-                    with session.begin_nested():
+                    with self.db.begin_nested():
                         loan.end = utc_now()
                 except StaleDataError:
                     # This can happen if this callback happened while we were returning this

diff --git a/tests/manager/api/controller/test_odl_notify.py b/tests/manager/api/controller/test_odl_notify.py
@@ -1,11 +1,10 @@
-from unittest.mock import PropertyMock, create_autospec, patch
+from unittest.mock import PropertyMock, create_autospec
 
 import pytest
 from flask import Response
 from freezegun import freeze_time
 from sqlalchemy.orm.exc import StaleDataError
 
-from palace.manager.api.controller import odl_notification
 from palace.manager.api.controller.odl_notification import ODLNotificationController
 from palace.manager.api.odl.api import OPDS2WithODLApi
 from palace.manager.api.problem_details import (
@@ -266,14 +265,11 @@ def test__process_notification_already_deleted(
         mock_loan.license_pool.collection.integration_configuration.protocol = (
             db.protocol_string(Goals.LICENSE_GOAL, OPDS2WithODLApi)
         )
-        with (
-            flask_app_fixture.test_request_context(
-                "/",
-                method="POST",
-                library=odl_fixture.library,
-                data=odl_fixture.loan_status_document("revoked").model_dump_json(),
-            ),
-            patch.object(odl_notification, "object_session"),
+        with flask_app_fixture.test_request_context(
+            "/",
+            method="POST",
+            library=odl_fixture.library,
+            data=odl_fixture.loan_status_document("revoked").model_dump_json(),
         ):
             response = odl_fixture.controller._process_notification(mock_loan)
         assert response.status_code == 204