Spaces:
Sleeping
Sleeping
| import json | |
| import os | |
| import threading | |
| import unittest | |
| from unittest.mock import patch | |
| from urllib.error import HTTPError | |
| from urllib.request import Request, urlopen | |
| from app import CONFIG, DemoServer, demo_address, validate_request | |
| class FakeEngine: | |
| def __init__(self): | |
| self.lock = threading.Lock() | |
| self.last = None | |
| def healthy(self): | |
| return True | |
| def decide(self, body): | |
| self.last = body | |
| return {"usage": {"oracle_calls": 0}, "decision": {"task_label": "card_arrival"}} | |
| class GatewayTests(unittest.TestCase): | |
| def setUpClass(cls): | |
| cls.engine = FakeEngine() | |
| cls.server = DemoServer(("127.0.0.1", 0), cls.engine) | |
| cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True) | |
| cls.thread.start() | |
| cls.url = "http://127.0.0.1:%d" % cls.server.server_port | |
| def tearDownClass(cls): | |
| cls.server.shutdown() | |
| cls.server.server_close() | |
| cls.thread.join(timeout=3) | |
| def request(self, path, body=None, content_type="application/json", raw=None): | |
| data = raw if raw is not None else None if body is None else json.dumps(body).encode() | |
| req = Request(self.url + path, data=data, headers={"Content-Type": content_type}) | |
| try: | |
| response = urlopen(req, timeout=5) | |
| except HTTPError as error: | |
| response = error | |
| with response: | |
| return response.status, response.read(), response.headers | |
| def test_valid_decision_forces_safe_options(self): | |
| status, body, _ = self.request("/api/decide", {"text": " hello ", "skill": "banking77"}) | |
| self.assertEqual(status, 200) | |
| self.assertEqual(json.loads(body)["usage"]["oracle_calls"], 0) | |
| self.assertEqual(self.engine.last["input"], {"text": "hello"}) | |
| self.assertFalse(self.engine.last["options"]["allow_oracle"]) | |
| def test_all_skills_and_profiles(self): | |
| for skill in ("banking77", "clinc150", "massive"): | |
| for profile in ("balanced", "quality-first", "cost-saver"): | |
| self.assertEqual(self.request("/api/decide", {"text": "hello", "skill": skill, "profile": profile})[0], 200) | |
| def test_invalid_inputs_rejected(self): | |
| for body in [None, [], "hello", {}, {"text": "", "skill": "banking77"}, | |
| {"text": "x" * 2001, "skill": "banking77"}, | |
| {"text": "hello", "skill": "../../admin"}, | |
| {"text": "hello", "skill": ["banking77"]}, | |
| {"text": "hello", "skill": "banking77", "profile": "unknown"}, | |
| {"text": "hello", "skill": "banking77", "oracle": True}]: | |
| with self.subTest(body=str(body)[:60]): | |
| status, _, _ = self.request("/api/decide", raw=json.dumps(body).encode()) | |
| self.assertEqual(status, 400) | |
| def test_body_limits(self): | |
| self.assertEqual(self.request("/api/decide", raw=b"x" * 16385)[0], 413) | |
| self.assertEqual(self.request("/api/decide", raw=b"{")[0], 400) | |
| self.assertEqual(self.request("/api/decide", raw=b"\xff")[0], 400) | |
| self.assertEqual(self.request("/api/decide", raw=b"{}", content_type="text/plain")[0], 415) | |
| def test_busy_returns_retryable_error(self): | |
| with self.engine.lock: | |
| self.assertEqual(self.request("/api/decide", {"text": "hello", "skill": "banking77"})[0], 429) | |
| def test_no_admin_or_filesystem_proxy(self): | |
| for path in ["/v1/admin/keys", "/v1/route", "/../app.py", "/app.py", "/model.cmf", "/config.json"]: | |
| self.assertEqual(self.request(path)[0], 404) | |
| self.assertEqual(self.request(path, {})[0], 404) | |
| def test_health_and_static_security_headers(self): | |
| for path in ["/", "/style.css", "/app.js", "/healthz"]: | |
| status, body, headers = self.request(path) | |
| self.assertEqual(status, 200) | |
| self.assertTrue(body) | |
| self.assertEqual(headers["X-Content-Type-Options"], "nosniff") | |
| self.assertIn("script-src 'self'", headers["Content-Security-Policy"]) | |
| self.assertFalse(json.loads(self.request("/healthz")[1])["oracle"]) | |
| def test_configuration_disables_learning_oracle_cache(self): | |
| for name in ("oracle", "learning", "cache"): | |
| self.assertFalse(CONFIG[name]["enabled"]) | |
| def test_hf_hostname_does_not_override_bind_address(self): | |
| with patch.dict(os.environ, {"SPACE_HOST": "infosave-cmf-decision.hf.space"}, clear=True): | |
| self.assertEqual(demo_address(), ("127.0.0.1", 7860)) | |
| with patch.dict(os.environ, {"SPACE_HOST": "infosave-cmf-decision.hf.space", "CMF_DEMO_HOST": "0.0.0.0", "PORT": "7860"}, clear=True): | |
| self.assertEqual(demo_address(), ("0.0.0.0", 7860)) | |
| def test_text_is_data_not_a_command(self): | |
| text = "$(touch /tmp/not-executed); <script>alert(1)</script>" | |
| self.assertEqual(validate_request({"text": text, "skill": "banking77"})["input"]["text"], text) | |
| if __name__ == "__main__": | |
| unittest.main() | |