xas-test-library/tests/test_console_outputs.py

from http.server import BaseHTTPRequestHandler, HTTPServer
import json
from pathlib import Path
import tempfile
import threading
import unittest
import urllib.error
import urllib.request
from console.artifacts import Collector
from console.job_main import BudgetRelay
from console.settings import Settings


class OutputTests(unittest.TestCase):
    def test_final_file_log_drain_includes_fast_script_tail_larger_than_one_chunk(self):
        with tempfile.TemporaryDirectory() as temp:
            root = Path(temp)
            logs, run = root / "logs", root / "run"
            logs.mkdir()
            run.mkdir()
            collector = Collector(Settings(root, []), run, [str(logs)])
            payload = b"file log\n" * 10000
            (logs / "new.log").write_bytes(payload)
            parts = []
            collector.finish_follow(parts.append)
            self.assertEqual(b"".join(parts), "\n[文件日志 new.log]\n".encode() + payload)
            self.assertEqual(collector.follow(parts.append), 0)

    def test_proxy_does_not_redirect_credentials_or_add_requests(self):
        seen = []
        class Redirect(BaseHTTPRequestHandler):
            def log_message(self, *args):
                pass
            def do_GET(self):
                seen.append(self.path)
                self.send_response(302)
                self.send_header("Location", "/unexpected")
                self.end_headers()
        upstream = HTTPServer(("127.0.0.1", 0), Redirect)
        relay = BudgetRelay(("127.0.0.1", 0), f"http://127.0.0.1:{upstream.server_port}", "test-only-credential")
        threads = [threading.Thread(target=s.serve_forever, daemon=True) for s in (upstream, relay)]
        for thread in threads:
            thread.start()
        try:
            with self.assertRaises(urllib.error.HTTPError) as rejected:
                urllib.request.urlopen(f"http://127.0.0.1:{relay.server_port}/v1/models", timeout=5)
            self.assertEqual(rejected.exception.code, 302)
            self.assertEqual(seen, ["/v1/models"])
            self.assertEqual(relay.gets, 1)
        finally:
            for server in (relay, upstream):
                server.shutdown()
                server.server_close()
            for thread in threads:
                thread.join()

    def test_only_changed_external_outputs_are_archived_and_originals_preserved(self):
        with tempfile.TemporaryDirectory() as temp:
            root = Path(temp)
            output, run = root / "original_logs", root / "run"
            output.mkdir()
            run.mkdir()
            (run / "work").mkdir()
            old = output / "old.json"
            old.write_text('{"old":true}')
            log = output / "current.log"
            log.write_bytes(b"previous\n")
            collector = Collector(Settings(root, []), run, [str(output)])
            with log.open("ab") as handle:
                handle.write("新日志\n".encode())
            messages = []
            collector.follow(messages.append)
            collector.follow(messages.append)
            self.assertEqual(len(messages), 1)
            self.assertIn("新日志", messages[0].decode())
            self.assertNotIn("previous", messages[0].decode())
            (run / "work/result.json").write_text('{"ok":true}')
            items = collector.collect()
            self.assertEqual({i["name"] for i in items}, {"current.log", "result.json"})
            self.assertEqual(old.read_text(), '{"old":true}')
            self.assertTrue(log.exists())
            log.write_text("overwritten later")
            archived = next(i for i in items if i["name"] == "current.log")
            self.assertIn("新日志", (run / "artifacts" / archived["id"]).read_text(encoding="utf-8"))

    def test_model_proxy_budget_and_auth_are_per_run_without_retry(self):
        requests = []
        class Upstream(BaseHTTPRequestHandler):
            def log_message(self, *args):
                pass
            def do_GET(self):
                requests.append((self.path, self.headers.get("Authorization")))
                self.send_response(200)
                self.end_headers()
                self.wfile.write(b'{"data":[]}')
            def do_POST(self):
                self.rfile.read(int(self.headers["Content-Length"]))
                requests.append((self.path, self.headers.get("Authorization")))
                self.send_response(503)
                self.end_headers()
        upstream = HTTPServer(("127.0.0.1", 0), Upstream)
        relay = BudgetRelay(("127.0.0.1", 0), f"http://127.0.0.1:{upstream.server_port}", "test-only-credential")
        threads = [threading.Thread(target=s.serve_forever, daemon=True) for s in (upstream, relay)]
        for t in threads:
            t.start()
        url = f"http://127.0.0.1:{relay.server_port}"
        try:
            with urllib.request.urlopen(url + "/v1/models", timeout=5) as response:
                self.assertEqual(response.status, 200)
            with self.assertRaises(urllib.error.HTTPError) as rejected:
                urllib.request.urlopen(url + "/v1/models", timeout=5)
            self.assertEqual(rejected.exception.code, 403)
            data = json.dumps({"model": "deepseek-v4.1-flash", "max_tokens": 512}).encode()
            for i in range(4):
                with self.assertRaises(urllib.error.HTTPError) as rejected:
                    urllib.request.urlopen(url + "/v1/chat/completions", data, timeout=5)
                self.assertEqual(rejected.exception.code, 503 if i < 3 else 403)
            self.assertEqual(len(requests), 4)
            self.assertTrue(all(auth == "Bearer test-only-credential" for _, auth in requests))
            self.assertNotIn("test-only-credential", json.dumps(relay.records))
        finally:
            for server in (relay, upstream):
                server.shutdown()
                server.server_close()
            for thread in threads:
                thread.join()