refactor: clean up async tests and remove redundant comments

Author: taeold

src/functions_framework/aio/__init__.py

@@ -58,7 +58,6 @@
 
 
 async def _crash_handler(request, exc):
-    # Log the exception
     logger = logging.getLogger()
     tb_lines = traceback.format_exception(type(exc), exc, exc.__traceback__)
     tb_text = "".join(tb_lines)
@@ -72,7 +71,6 @@ async def _crash_handler(request, exc):
         log_entry = {"message": error_msg, "levelname": "ERROR"}
         logger.error(json.dumps(log_entry))
     else:
-        # Execution ID logging not enabled, log plain text
         logger.error(error_msg)
 
     headers = {_FUNCTION_STATUS_HEADER_FIELD: _CRASH}
@@ -191,19 +189,16 @@ def _enable_execution_id_logging():
 
 
 def _configure_app_execution_id_logging():
-    # Logging needs to be configured before app logger is accessed
     import logging
     import logging.config
 
-    # Configure root logger to use our custom handler
     root_logger = logging.getLogger()
     root_logger.setLevel(logging.INFO)
 
     # Remove existing handlers
     for handler in root_logger.handlers[:]:
         root_logger.removeHandler(handler)
 
-    # Add our custom handler that adds execution ID
     handler = logging.StreamHandler(
         execution_id.LoggingHandlerAddExecutionId(sys.stderr)
     )

src/functions_framework/execution_id.py

@@ -51,11 +51,9 @@ def __init__(self, execution_id=None, span_id=None):
 
 
 def _get_current_context():
-    # First try to get from async context
     context = execution_context_var.get()
     if context is not None:
         return context
-    # Fall back to Flask context for sync
     return (  # pragma: no cover
         flask.g.execution_id_context
         if flask.has_request_context() and "execution_id_context" in flask.g
@@ -64,8 +62,6 @@ def _get_current_context():
 
 
 def _set_current_context(context):
-    # Set in both contexts to support both sync and async
-    # Set in contextvars for async
     execution_context_var.set(context)
     # Also set in Flask context if available for sync
     if flask.has_request_context():
@@ -82,8 +78,7 @@ def _generate_execution_id():
 def _extract_context_from_headers(headers):
     """Extract execution context from request headers."""
     execution_id = headers.get(EXECUTION_ID_REQUEST_HEADER)
-
-    # Try to get span ID from trace context header
+    
     trace_context = re.match(
         _TRACE_CONTEXT_REGEX_PATTERN,
         headers.get(TRACE_CONTEXT_REQUEST_HEADER, ""),
@@ -113,7 +108,6 @@ def __init__(self, app):
 
     async def __call__(self, scope, receive, send):
         if scope["type"] == "http":
-            # Extract existing execution ID or generate a new one
             execution_id_header = b"function-execution-id"
             trace_context_header = b"x-cloud-trace-context"
             execution_id = None
@@ -127,22 +121,18 @@ async def __call__(self, scope, receive, send):
 
             if not execution_id:
                 execution_id = _generate_execution_id()
-                # Add the execution ID to headers
                 new_headers = list(scope.get("headers", []))
                 new_headers.append(
                     (execution_id_header, execution_id.encode("latin-1"))
                 )
                 scope["headers"] = new_headers
 
-            # Store execution context in ASGI scope for recovery in case of context loss
-            # Parse trace context to extract span ID
             span_id = None
             if trace_context:
                 trace_match = re.match(_TRACE_CONTEXT_REGEX_PATTERN, trace_context)
                 if trace_match:
                     span_id = trace_match.group("span_id")
 
-            # Store in scope for potential recovery
             scope["execution_context"] = ExecutionContext(execution_id, span_id)
 
         await self.app(scope, receive, send)  # pragma: no cover
@@ -169,9 +159,7 @@ def wrapper(*args, **kwargs):
 
             with stderr_redirect, stdout_redirect:
                 result = view_function(*args, **kwargs)
-
                 # Context cleanup happens automatically via Flask's request context
-                # No need to manually clean up flask.g
                 return result
 
         return wrapper
@@ -195,14 +183,10 @@ def set_execution_context_async(enable_id_logging=False):
     def decorator(view_function):
         @functools.wraps(view_function)
         async def async_wrapper(request, *args, **kwargs):
-            # Extract execution context from headers
             context = _extract_context_from_headers(request.headers)
-
-            # Set context using contextvars
             token = execution_context_var.set(context)
 
             with stderr_redirect, stdout_redirect:
-                # Handle both sync and async functions
                 if inspect.iscoroutinefunction(view_function):
                     result = await view_function(request, *args, **kwargs)
                 else:
@@ -215,10 +199,7 @@ async def async_wrapper(request, *args, **kwargs):
 
         @functools.wraps(view_function)
         def sync_wrapper(request, *args, **kwargs):  # pragma: no cover
-            # For sync functions, we still need to set up the context
             context = _extract_context_from_headers(request.headers)
-
-            # Set context using contextvars
             token = execution_context_var.set(context)
 
             with stderr_redirect, stdout_redirect:
@@ -229,7 +210,6 @@ def sync_wrapper(request, *args, **kwargs):  # pragma: no cover
                 execution_context_var.reset(token)
                 return result
 
-        # Return appropriate wrapper based on whether the function is async
         if inspect.iscoroutinefunction(view_function):
             return async_wrapper
         else:

tests/test_execution_id_async.py

@@ -11,13 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import asyncio
 import json
 import pathlib
-import re
-import sys
-
-from functools import partial
 from unittest.mock import Mock
 
 import pytest
@@ -32,8 +27,7 @@
 TEST_SPAN_ID = "123456"
 
 
-@pytest.mark.asyncio
-async def test_async_user_function_can_retrieve_execution_id_from_header():
+def test_async_user_function_can_retrieve_execution_id_from_header():
     source = TEST_FUNCTIONS_DIR / "execution_id" / "async_main.py"
     target = "async_function"
     app = create_asgi_app(target, source)
@@ -49,8 +43,7 @@ async def test_async_user_function_can_retrieve_execution_id_from_header():
     assert resp.json()["execution_id"] == TEST_EXECUTION_ID
 
 
-@pytest.mark.asyncio
-async def test_async_uncaught_exception_in_user_function_sets_execution_id(
+def test_async_uncaught_exception_in_user_function_sets_execution_id(
     capsys, monkeypatch
 ):
     monkeypatch.setenv("LOG_EXECUTION_ID", "true")
@@ -71,8 +64,7 @@ async def test_async_uncaught_exception_in_user_function_sets_execution_id(
     assert f'"execution_id": "{TEST_EXECUTION_ID}"' in record.err
 
 
-@pytest.mark.asyncio
-async def test_async_print_from_user_function_sets_execution_id(capsys, monkeypatch):
+def test_async_print_from_user_function_sets_execution_id(capsys, monkeypatch):
     monkeypatch.setenv("LOG_EXECUTION_ID", "true")
     source = TEST_FUNCTIONS_DIR / "execution_id" / "async_main.py"
     target = "async_print_message"
@@ -91,8 +83,7 @@ async def test_async_print_from_user_function_sets_execution_id(capsys, monkeypa
     assert '"message": "some-message"' in record.out
 
 
-@pytest.mark.asyncio
-async def test_async_log_from_user_function_sets_execution_id(capsys, monkeypatch):
+def test_async_log_from_user_function_sets_execution_id(capsys, monkeypatch):
     monkeypatch.setenv("LOG_EXECUTION_ID", "true")
     source = TEST_FUNCTIONS_DIR / "execution_id" / "async_main.py"
     target = "async_log_message"
@@ -111,8 +102,7 @@ async def test_async_log_from_user_function_sets_execution_id(capsys, monkeypatc
     assert '"custom-field": "some-message"' in record.err
 
 
-@pytest.mark.asyncio
-async def test_async_user_function_can_retrieve_generated_execution_id(monkeypatch):
+def test_async_user_function_can_retrieve_generated_execution_id(monkeypatch):
     monkeypatch.setattr(
         execution_id, "_generate_execution_id", lambda: TEST_EXECUTION_ID
     )
@@ -130,8 +120,7 @@ async def test_async_user_function_can_retrieve_generated_execution_id(monkeypat
     assert resp.json()["execution_id"] == TEST_EXECUTION_ID
 
 
-@pytest.mark.asyncio
-async def test_async_does_not_set_execution_id_when_not_enabled(capsys):
+def test_async_does_not_set_execution_id_when_not_enabled(capsys):
     source = TEST_FUNCTIONS_DIR / "execution_id" / "async_main.py"
     target = "async_print_message"
     app = create_asgi_app(target, source)
@@ -149,8 +138,7 @@ async def test_async_does_not_set_execution_id_when_not_enabled(capsys):
     assert "some-message" in record.out
 
 
-@pytest.mark.asyncio
-async def test_async_concurrent_requests_maintain_separate_execution_ids(
+def test_async_concurrent_requests_maintain_separate_execution_ids(
     capsys, monkeypatch
 ):
     monkeypatch.setenv("LOG_EXECUTION_ID", "true")
@@ -236,8 +224,7 @@ def make_request(client, message, exec_id):
         ),
     ],
 )
-@pytest.mark.asyncio
-async def test_async_set_execution_context_headers(
+def test_async_set_execution_context_headers(
     headers, expected_execution_id, expected_span_id, should_generate
 ):
     source = TEST_FUNCTIONS_DIR / "execution_id" / "async_main.py"
@@ -260,8 +247,6 @@ async def test_async_set_execution_context_headers(
 @pytest.mark.asyncio
 async def test_crash_handler_without_context_sets_execution_id():
     """Test that crash handler returns proper error response with crash header."""
-    from unittest.mock import Mock
-
     from functions_framework.aio import _crash_handler
 
     # Create a mock request
@@ -281,8 +266,7 @@ async def test_crash_handler_without_context_sets_execution_id():
     assert response.headers["X-Google-Status"] == "crash"
 
 
-@pytest.mark.asyncio
-async def test_async_decorator_with_sync_function():
+def test_async_decorator_with_sync_function():
     """Test that the async decorator handles sync functions properly."""
     from functions_framework.execution_id import set_execution_context_async