Files
huni-probe-agent/test_main.py
huni f18eeaf334 harden probe agent v1.0.3
安全鉴权、默认只读、显式提交开关、状态接口与测试
2026-08-26 17:16:02 +08:00

117 lines
4.8 KiB
Python

import unittest
from unittest import mock
import main
class TokenResolutionTests(unittest.TestCase):
def test_placeholder_legacy_token_is_rejected(self):
token, source, error = main._resolve_token({"EXTERNAL_SERVICE_TOKEN": "tmp"})
self.assertEqual(token, "")
self.assertEqual(source, "EXTERNAL_SERVICE_TOKEN")
self.assertEqual(error, "token_missing_or_placeholder")
def test_token_file_has_precedence(self):
token, source, error = main._resolve_token(
{"XC_TOKEN_FILE": "/secret/token", "XC_TOKEN": "env-token"},
read_text=lambda _path: "file-token\n",
)
self.assertEqual(token, "file-token")
self.assertEqual(source, "XC_TOKEN_FILE")
self.assertIsNone(error)
def test_unreadable_explicit_token_file_does_not_fall_back(self):
def fail(_path):
raise OSError("not mounted")
token, source, error = main._resolve_token(
{"XC_TOKEN_FILE": "/secret/token", "XC_TOKEN": "env-token"},
read_text=fail,
)
self.assertEqual(token, "")
self.assertEqual(source, "XC_TOKEN_FILE")
self.assertEqual(error, "token_file_unreadable")
class ProbeTests(unittest.TestCase):
def setUp(self):
main._reset_state()
def test_missing_token_makes_no_request(self):
with mock.patch.object(main, "STRATEGY_ID", "strategy-1"), mock.patch.object(
main, "_resolve_token", return_value=("", "EXTERNAL_SERVICE_TOKEN", "token_missing_or_placeholder")
), mock.patch.object(main, "_req") as request:
main.probe()
request.assert_not_called()
self.assertEqual(main._snapshot()["phase"], "blocked_missing_auth")
def test_read_only_mode_never_calls_post(self):
calls = []
def request(method, path, _token, body=None, timeout=30):
calls.append((method, path, body, timeout))
return {"http": 200, "code": 0, "data": {}}
with mock.patch.object(main, "STRATEGY_ID", "strategy-1"), mock.patch.object(
main, "ALLOW_SUBMIT", False
), mock.patch.object(main, "_resolve_token", return_value=("valid-token", "XC_TOKEN", None)), mock.patch.object(
main, "_req", side_effect=request
):
main.probe()
self.assertEqual([call[0] for call in calls], ["GET"])
state = main._snapshot()
self.assertEqual(state["phase"], "read_only_complete")
self.assertTrue(state["auth_ready"])
def test_failed_auth_never_calls_post(self):
with mock.patch.object(main, "STRATEGY_ID", "strategy-1"), mock.patch.object(
main, "ALLOW_SUBMIT", True
), mock.patch.object(main, "_resolve_token", return_value=("valid-token", "XC_TOKEN", None)), mock.patch.object(
main, "_req", return_value={"http": 401, "code": 40100, "message": "unauthorized"}
) as request:
main.probe()
self.assertEqual(request.call_count, 1)
self.assertEqual(request.call_args.args[0], "GET")
self.assertEqual(main._snapshot()["phase"], "auth_failed")
def test_submit_mode_calls_task_add_once(self):
responses = [
{"http": 200, "code": 0, "data": {}},
{"http": 200, "code": 0, "data": "framework: vllm\n"},
{"http": 200, "code": 0, "data": {"taskId": "task-1"}},
]
with mock.patch.object(main, "STRATEGY_ID", "strategy-1"), mock.patch.object(
main, "ALLOW_SUBMIT", True
), mock.patch.object(main, "_resolve_token", return_value=("valid-token", "XC_TOKEN", None)), mock.patch.object(
main, "_req", side_effect=responses
) as request:
main.probe()
methods_and_paths = [(call.args[0], call.args[1]) for call in request.call_args_list]
self.assertEqual(sum(path == "/api/adapt/task/add" for _, path in methods_and_paths), 1)
state = main._snapshot()
self.assertEqual(state["phase"], "submit_complete")
self.assertEqual(state["results"]["task_add"]["task_id"], "task-1")
def test_failed_build_config_prevents_task_add(self):
responses = [
{"http": 200, "code": 0, "data": {}},
{"http": 500, "code": 50000, "message": "build failed"},
]
with mock.patch.object(main, "STRATEGY_ID", "strategy-1"), mock.patch.object(
main, "ALLOW_SUBMIT", True
), mock.patch.object(main, "_resolve_token", return_value=("valid-token", "XC_TOKEN", None)), mock.patch.object(
main, "_req", side_effect=responses
) as request:
main.probe()
self.assertEqual(request.call_count, 2)
self.assertNotIn("/api/adapt/task/add", [call.args[1] for call in request.call_args_list])
self.assertEqual(main._snapshot()["phase"], "build_config_failed")
if __name__ == "__main__":
unittest.main()