From e826ab35eddce35f9a7205a5c894c80add76f6dc Mon Sep 17 00:00:00 2001
From: Jonathan Green <jonathan@razzard.com>
Date: Mon, 21 Oct 2024 15:45:05 -0300
Subject: [PATCH] Mypy fix

---
 .../manager/api/controller/odl_notification.py   |  5 ++---
 tests/manager/api/controller/test_odl_notify.py  | 16 ++++++----------
 2 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/src/palace/manager/api/controller/odl_notification.py b/src/palace/manager/api/controller/odl_notification.py
index 89012c4b4..3e45ff4de 100644
--- 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
index 3aca95f6b..f02433b1a 100644
--- 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