Refactor tests/connections/test_pybricks.py to use the same ExitStack pattern we just applied to test_cli.py.

Author: 1e0ng

tests/connections/test_pybricks.py

@@ -1,8 +1,10 @@
 """Tests for the pybricks connection module."""
 
 import asyncio
+import contextlib
+import os
 import tempfile
-from unittest.mock import AsyncMock, PropertyMock
+from unittest.mock import AsyncMock, PropertyMock, patch
 
 import pytest
 from reactivex.subject import Subject
@@ -32,16 +34,17 @@ async def test_download_modern_protocol(self):
         )
         hub._capability_flags = HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6
 
-        with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as temp:
+        with contextlib.ExitStack() as stack:
+            # Create and manage temporary file
+            temp = stack.enter_context(
+                tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False)
+            )
             temp.write("print('test')")
             temp_path = temp.name
-        try:
+            stack.callback(os.unlink, temp_path)
+
             await hub.download(temp_path)
             hub.download_user_program.assert_called_once()
-        finally:
-            import os
-
-            os.unlink(temp_path)
 
     @pytest.mark.asyncio
     async def test_download_legacy_firmware(self):
@@ -56,16 +59,17 @@ async def test_download_legacy_firmware(self):
         )
         hub._capability_flags = HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6
 
-        with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as temp:
+        with contextlib.ExitStack() as stack:
+            # Create and manage temporary file
+            temp = stack.enter_context(
+                tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False)
+            )
             temp.write("print('test')")
             temp_path = temp.name
-        try:
+            stack.callback(os.unlink, temp_path)
+
             await hub.download(temp_path)
             hub.download_user_program.assert_called_once()
-        finally:
-            import os
-
-            os.unlink(temp_path)
 
     @pytest.mark.asyncio
     async def test_download_unsupported_capabilities(self):
@@ -79,19 +83,20 @@ async def test_download_unsupported_capabilities(self):
         )
         hub._capability_flags = 0
 
-        with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as temp:
+        with contextlib.ExitStack() as stack:
+            # Create and manage temporary file
+            temp = stack.enter_context(
+                tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False)
+            )
             temp.write("print('test')")
             temp_path = temp.name
-        try:
+            stack.callback(os.unlink, temp_path)
+
             with pytest.raises(
                 RuntimeError,
                 match="Hub is not compatible with any of the supported file formats",
             ):
                 await hub.download(temp_path)
-        finally:
-            import os
-
-            os.unlink(temp_path)
 
     @pytest.mark.asyncio
     async def test_download_compile_error(self):
@@ -104,17 +109,27 @@ async def test_download_compile_error(self):
             return_value=ConnectionState.CONNECTED
         )
         hub._capability_flags = HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6
+        hub._max_user_program_size = 1000  # Set a reasonable size limit
 
-        with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as temp:
+        with contextlib.ExitStack() as stack:
+            # Create and manage temporary file
+            temp = stack.enter_context(
+                tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False)
+            )
             temp.write("print('test'  # Missing closing parenthesis")
             temp_path = temp.name
-        try:
-            with pytest.raises(SyntaxError):
-                await hub.download(temp_path)
-        finally:
-            import os
+            stack.callback(os.unlink, temp_path)
+
+            # Mock compile_multi_file to raise SyntaxError
+            stack.enter_context(
+                patch(
+                    "pybricksdev.connections.pybricks.compile_multi_file",
+                    side_effect=SyntaxError("invalid syntax"),
+                )
+            )
 
-            os.unlink(temp_path)
+            with pytest.raises(SyntaxError, match="invalid syntax"):
+                await hub.download(temp_path)
 
     @pytest.mark.asyncio
     async def test_run_modern_protocol(self):
@@ -139,10 +154,15 @@ async def test_run_modern_protocol(self):
         hub._stdout_line_queue = asyncio.Queue()
         hub._enable_line_handler = True
 
-        with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as temp:
+        with contextlib.ExitStack() as stack:
+            # Create and manage temporary file
+            temp = stack.enter_context(
+                tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False)
+            )
             temp.write("print('test')")
             temp_path = temp.name
-        try:
+            stack.callback(os.unlink, temp_path)
+
             # Start the run task
             run_task = asyncio.create_task(hub.run(temp_path))
 
@@ -160,8 +180,3 @@ async def test_run_modern_protocol(self):
             # Verify the expected calls were made
             hub.download_user_program.assert_called_once()
             hub.start_user_program.assert_called_once()
-        finally:
-            import os
-
-            os.unlink(temp_path)
-