fix: refactor tests

Author: taeold

tests/test_cli.py

@@ -128,9 +128,17 @@ def test_asgi_cli(monkeypatch):
     assert asgi_server.run.calls == [pretend.call("0.0.0.0", 8080)]
 
 
-def test_auto_asgi_for_aio_decorated_functions(monkeypatch):
-    original_asgi_functions = _function_registry.ASGI_FUNCTIONS.copy()
+@pytest.fixture
+def clean_registry():
+    """Save and restore function registry state."""
+    original_asgi = _function_registry.ASGI_FUNCTIONS.copy()
     _function_registry.ASGI_FUNCTIONS.clear()
+    yield
+    _function_registry.ASGI_FUNCTIONS.clear()
+    _function_registry.ASGI_FUNCTIONS.update(original_asgi)
+
+
+def test_auto_asgi_for_aio_decorated_functions(monkeypatch, clean_registry):
     _function_registry.ASGI_FUNCTIONS.add("my_aio_func")
 
     asgi_app = pretend.stub()
@@ -148,13 +156,8 @@ def test_auto_asgi_for_aio_decorated_functions(monkeypatch):
     assert create_asgi_app.calls == [pretend.call("my_aio_func", None, "http")]
     assert asgi_server.run.calls == [pretend.call("0.0.0.0", 8080)]
 
-    _function_registry.ASGI_FUNCTIONS.clear()
-    _function_registry.ASGI_FUNCTIONS.update(original_asgi_functions)
-
 
-def test_no_auto_asgi_for_regular_functions(monkeypatch):
-    original_asgi_functions = _function_registry.ASGI_FUNCTIONS.copy()
-    _function_registry.ASGI_FUNCTIONS.clear()
+def test_no_auto_asgi_for_regular_functions(monkeypatch, clean_registry):
 
     app = pretend.stub()
     create_app = pretend.call_recorder(lambda *a, **k: app)
@@ -169,6 +172,3 @@ def test_no_auto_asgi_for_regular_functions(monkeypatch):
 
     assert create_app.calls == [pretend.call("regular_func", None, "http")]
     assert flask_server.run.calls == [pretend.call("0.0.0.0", 8080)]
-
-    _function_registry.ASGI_FUNCTIONS.clear()
-    _function_registry.ASGI_FUNCTIONS.update(original_asgi_functions)

tests/test_decorator_functions.py

@@ -132,13 +132,24 @@ def test_aio_http_dict_response():
     assert resp.json() == {"message": "hello", "count": 42, "success": True}
 
 
-def test_aio_decorators_register_asgi_functions():
-    """Test that @aio decorators add function names to ASGI_FUNCTIONS registry."""
+@pytest.fixture
+def clean_registry():
+    """Save and restore registry state."""
     original_registry_map = registry.REGISTRY_MAP.copy()
     original_asgi_functions = registry.ASGI_FUNCTIONS.copy()
     registry.REGISTRY_MAP.clear()
     registry.ASGI_FUNCTIONS.clear()
 
+    yield
+
+    registry.REGISTRY_MAP.clear()
+    registry.REGISTRY_MAP.update(original_registry_map)
+    registry.ASGI_FUNCTIONS.clear()
+    registry.ASGI_FUNCTIONS.update(original_asgi_functions)
+
+
+def test_aio_decorators_register_asgi_functions(clean_registry):
+    """Test that @aio decorators add function names to ASGI_FUNCTIONS registry."""
     from functions_framework.aio import cloud_event, http
 
     @http
@@ -165,8 +176,3 @@ def test_cloud_event_sync(event):
 
     assert "test_http_sync" in registry.ASGI_FUNCTIONS
     assert "test_cloud_event_sync" in registry.ASGI_FUNCTIONS
-
-    registry.REGISTRY_MAP.clear()
-    registry.REGISTRY_MAP.update(original_registry_map)
-    registry.ASGI_FUNCTIONS.clear()
-    registry.ASGI_FUNCTIONS.update(original_asgi_functions)