From b9ecb27560f06e9e1fe6e3bad8fa350d0c201cb3 Mon Sep 17 00:00:00 2001
From: huanghuoguoguo <1051233107@qq.com>
Date: Sun, 17 May 2026 10:41:02 +0800
Subject: [PATCH] test: address agent runner review comments

---
 tests/unit_tests/agent/test_handler_auth.py   | 33 +++++--------------
 .../unit_tests/agent/test_session_registry.py | 10 +++---
 2 files changed, 13 insertions(+), 30 deletions(-)

diff --git a/tests/unit_tests/agent/test_handler_auth.py b/tests/unit_tests/agent/test_handler_auth.py
index 3771c992..f4f303c6 100644
--- a/tests/unit_tests/agent/test_handler_auth.py
+++ b/tests/unit_tests/agent/test_handler_auth.py
@@ -431,11 +431,9 @@ class TestAuthorizationPathDifferentiation:
         # Regular plugin call has no run_id
         run_id = None
 
-        # Authorization check should be skipped when run_id is None
+        # Authorization check should be skipped when run_id is None.
         # This is handled in handler.py with: if run_id: ...
-        if run_id:
-            # This block should NOT execute for regular plugin calls
-            raise AssertionError('Authorization check should not run for regular plugin calls')
+        assert run_id is None
 
         # For regular plugins:
         # - INVOKE_LLM: unrestricted access to any model
@@ -661,12 +659,7 @@ class TestHandlerActionAuthorization:
         run_id = None
 
         # In handler.py, authorization check is inside: if run_id: ...
-        # So when run_id is None, authorization is skipped
-        if run_id:
-            # This block should NOT execute
-            raise AssertionError('Should not execute authorization for no run_id')
-
-        # No authorization check - unrestricted access
+        # So when run_id is None, authorization is skipped.
         assert run_id is None
 
     @pytest.mark.asyncio
@@ -733,9 +726,6 @@ class TestHandlerActionAuthorization:
         run_id = None
 
         # Authorization check is inside: if run_id: ...
-        if run_id:
-            raise AssertionError('Should not execute authorization for no run_id')
-
         assert run_id is None
 
     @pytest.mark.asyncio
@@ -880,10 +870,8 @@ class TestNoRunIdBackwardCompatPath:
         run_id = None
         llm_model_uuid = 'any_model'
 
-        # Simulate handler logic
-        if run_id:
-            # This should NOT execute
-            raise AssertionError('Authorization should not run')
+        # Simulate handler logic: no run_id skips authorization.
+        assert run_id is None
 
         # Model can be any UUID (unrestricted)
         assert llm_model_uuid == 'any_model'
@@ -895,8 +883,7 @@ class TestNoRunIdBackwardCompatPath:
         tool_name = 'any_tool'
 
         # Handler.py line 463: if run_id: ...
-        if run_id:
-            raise AssertionError('Authorization should not run')
+        assert run_id is None
 
         assert tool_name == 'any_tool'
 
@@ -1982,9 +1969,8 @@ class TestBackwardCompatStorageNoRunId:
         # When run_id is None, validation is skipped
         run_id = None
 
-        # Simulate handler logic
-        if run_id:
-            raise AssertionError('Should not execute validation')
+        # Simulate handler logic: no run_id skips validation.
+        assert run_id is None
 
         # Storage access unrestricted for regular plugins
         assert run_id is None
@@ -1993,8 +1979,7 @@ class TestBackwardCompatStorageNoRunId:
         """GET_CONFIG_FILE without run_id skips Host-side validation."""
         run_id = None
 
-        if run_id:
-            raise AssertionError('Should not execute validation')
+        assert run_id is None
 
         # File access unrestricted for regular plugins
         assert run_id is None
diff --git a/tests/unit_tests/agent/test_session_registry.py b/tests/unit_tests/agent/test_session_registry.py
index 415c2188..c47205f8 100644
--- a/tests/unit_tests/agent/test_session_registry.py
+++ b/tests/unit_tests/agent/test_session_registry.py
@@ -27,8 +27,6 @@ class TestSessionRegistryBasic:
             models=[{'model_id': 'model_001', 'model_type': 'chat', 'provider': 'openai'}],
             tools=[{'tool_name': 'web_search', 'tool_type': 'builtin'}],
         )
-        session = make_session(run_id=run_id, resources=resources)
-
         await registry.register(
             run_id=run_id,
             runner_id='plugin:test/my-runner/default',
@@ -393,10 +391,10 @@ class TestGlobalRegistry:
         # In production, calling get_session_registry() multiple times
         # returns the same instance. We verify this by checking the
         # module-level variable directly.
-        import langbot.pkg.agent.runner.session_registry as registry_module
+        from langbot.pkg.agent.runner.session_registry import _global_registry
 
         # Check that the global variable exists and is either None or an instance
-        global_reg = registry_module._global_registry
+        global_reg = _global_registry
         if global_reg is None:
             # First call creates the instance
             registry1 = get_session_registry()
@@ -453,8 +451,8 @@ class TestThreadSafety:
             registry.get('run_1'),
         ]
 
-        results = await asyncio.gather(*tasks)
+        await asyncio.gather(*tasks)
 
         # After both complete, session should be unregistered
         result = await registry.get('run_1')
-        assert result is None
\ No newline at end of file
+        assert result is None