From a01790c2050c36584289d097cbe8261a322d4106 Mon Sep 17 00:00:00 2001
From: Brandur <brandur@brandur.org>
Date: Sat, 25 Jan 2025 14:35:17 -0800
Subject: [PATCH] Don't try to mock context manager; use a simple class instead

This one's related to #38. It turns out that while trying to mock a
context manager kind of works, it will do the wrong thing in edge cases
like when an exception is thrown from inside a `with` block, silently
swallowing it and causing a return that's completely wrong.

There may be some way to fix the mock to make it do the right thing, but
instead of getting fancier with these mocks that are already awful,
instead repair the problem by defining a plain class that implements
context manager and just use that.
---
 .github/workflows/ci.yaml |  4 ++--
 tests/client_test.py      | 21 ++++++++++++---------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 3135280..914f9b5 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -63,7 +63,7 @@ jobs:
         run: psql --echo-errors --quiet -c '\timing off' -c "CREATE DATABASE ${TEST_DATABASE_NAME};" ${ADMIN_DATABASE_URL}
 
       - name: river migrate-up
-        run: river migrate-up --database-url "$TEST_DATABASE_URL"
+        run: river migrate-up --database-url "$TEST_DATABASE_URL" --max-steps 5 # temporarily include max steps so tests can pass with unique fixes
 
       - name: Test
         run: rye test
@@ -109,7 +109,7 @@ jobs:
         run: psql --echo-errors --quiet -c '\timing off' -c "CREATE DATABASE ${DATABASE_NAME};" ${ADMIN_DATABASE_URL}
 
       - name: river migrate-up
-        run: river migrate-up --database-url "$DATABASE_URL"
+        run: river migrate-up --database-url "$DATABASE_URL" --max-steps 5 # temporarily include max steps so tests can pass with unique fixes
 
       - name: Run examples
         run: rye run python3 -m examples.all
diff --git a/tests/client_test.py b/tests/client_test.py
index b8e9723..37e53b0 100644
--- a/tests/client_test.py
+++ b/tests/client_test.py
@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -16,17 +16,20 @@ def mock_driver() -> DriverProtocol:
 
 @pytest.fixture
 def mock_exec(mock_driver) -> ExecutorProtocol:
-    def mock_context_manager(val) -> Mock:
-        context_manager_mock = MagicMock()
-        context_manager_mock.__enter__.return_value = val
-        context_manager_mock.__exit__.return_value = Mock()
-        return context_manager_mock
+    # Don't try to mock a context manager. It will cause endless pain around the
+    # edges like swallowing raised exceptions.
+    class TrivialContextManager:
+        def __init__(self, with_val):
+            self.with_val = with_val
 
-    # def mock_context_manager(val) -> Mock:
-    #     return Mock(__enter__=val, __exit__=Mock())
+        def __enter__(self):
+            return self.with_val
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            pass
 
     mock_exec = MagicMock(spec=ExecutorProtocol)
-    mock_driver.executor.return_value = mock_context_manager(mock_exec)
+    mock_driver.executor.return_value = TrivialContextManager(mock_exec)
 
     return mock_exec