codex-app-server/cloud/tests/test_resident.py

"""常驻监听任务单测（cloud/resident.py）：不跑真实模型、不联网，
抓取一律注入假 fetcher。覆盖创建校验、暂停保游标、恢复续跑、
结束终态、防重复轮询、仅变化建任务、失败退避、配额耗尽转暂停、
worker 驱动批量入口。"""
import http.server
import os
import socket
import sqlite3
import sys
import tempfile
import threading
import unittest
from pathlib import Path
from types import SimpleNamespace

try:
    import cloud.resident as resident_module
    from cloud.resident import (
        SSRFBlockedError,
        claim_due_watch,
        create_watch,
        default_fetch,
        end_watch,
        ensure_resident_tables,
        get_watch,
        list_polls,
        pause_watch,
        process_due_resident_watches,
        resume_watch,
        run_due_resident_watch,
        update_watch,
        validate_target_url,
    )
    from cloud.server import init_db
    from cloud.worker import Worker
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    import cloud.resident as resident_module
    from cloud.resident import (
        SSRFBlockedError,
        claim_due_watch,
        create_watch,
        default_fetch,
        end_watch,
        ensure_resident_tables,
        get_watch,
        list_polls,
        pause_watch,
        process_due_resident_watches,
        resume_watch,
        run_due_resident_watch,
        update_watch,
        validate_target_url,
    )
    from cloud.server import init_db
    from cloud.worker import Worker

URL = "https://example.com/data"


class _Fetcher:
    """可变假抓取：value 改了，下一次轮询就看到新内容。"""

    def __init__(self, value="v1"):
        self.value = value
        self.calls = []

    def __call__(self, url, timeout):
        self.calls.append((url, timeout))
        return self.value


def _failing_fetcher(url, timeout):
    raise RuntimeError("boom")


class ResidentTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        init_db(self.db_path)
        self.conn = sqlite3.connect(self.db_path)
        self.conn.row_factory = sqlite3.Row
        ensure_resident_tables(self.conn)
        self.now = 1700000000.0

    def tearDown(self):
        self.conn.close()
        self.tmp.cleanup()

    def _watch(self, name="watch1", **kw):
        kw.setdefault("interval_seconds", 60)
        return create_watch(self.conn, "proj", name, URL,
                            "变化时汇报", "owner1", now=self.now, **kw)

    def _tasks(self):
        cur = self.conn.execute("SELECT id, prompt, status FROM tasks")
        return cur.fetchall()

    def _set_quota(self, user_id, micros):
        self.conn.execute(
            "CREATE TABLE IF NOT EXISTS user_quotas("
            "user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, "
            " updated_at DOUBLE PRECISION)")
        self.conn.execute(
            "INSERT OR REPLACE INTO user_quotas(user_id,"
            " quota_usd_micros, updated_at) VALUES(?, ?, ?)",
            (user_id, micros, self.now))
        self.conn.commit()

    def test_create_validation(self):
        watch = self._watch()
        self.assertEqual(watch["status"], "running")
        self.assertTrue(watch["due"])
        self.assertEqual(watch["next_poll_at"], self.now)
        for kwargs, _ in (
                ({"target_url": "ftp://example.com/x"}, None),
                ({"target_url": "https://localhost/x"}, None),
                ({"target_url": "http://127.0.0.1/x"}, None),
                ({"target_url": "https://10.0.0.5/x"}, None),
                ({"target_url": "https://user:pw@example.com/x"}, None),
                ({"interval_seconds": 30}, None),
                ({"name": "watch1"}, None),
                ({"prompt": "  "}, None)):
            with self.assertRaises(ValueError):
                base = {"name": "other", "target_url": URL,
                        "prompt": "p", "interval_seconds": 60}
                base.update(kwargs)
                create_watch(self.conn, "proj", base["name"],
                             base["target_url"], base["prompt"], "owner1",
                             interval_seconds=base["interval_seconds"],
                             now=self.now)

    def test_baseline_unchanged_changed_only_change_creates_task(self):
        fetcher = _Fetcher("v1")
        watch = self._watch()
        poll = run_due_resident_watch(self.conn, watch["id"],
                                      fetcher=fetcher, now=self.now)
        self.assertEqual(poll["status"], "baseline")
        self.assertEqual(len(self._tasks()), 0)
        # 同一间隔内第二次：抢不到，不重复轮询。
        self.assertIsNone(run_due_resident_watch(
            self.conn, watch["id"], fetcher=fetcher, now=self.now))
        self.assertEqual(len(fetcher.calls), 1)
        self.now += 61
        poll = run_due_resident_watch(self.conn, watch["id"],
                                      fetcher=fetcher, now=self.now)
        self.assertEqual(poll["status"], "unchanged")
        self.assertEqual(len(self._tasks()), 0)
        fetcher.value = "v2"
        self.now += 61
        poll = run_due_resident_watch(self.conn, watch["id"],
                                      fetcher=fetcher, now=self.now)
        self.assertEqual(poll["status"], "changed")
        self.assertIsNotNone(poll["task_id"])
        tasks = self._tasks()
        self.assertEqual(len(tasks), 1)
        self.assertIn("v2", tasks[0][1])
        self.assertIn("变化时汇报", tasks[0][1])
        after = get_watch(self.conn, watch["id"], now=self.now)
        self.assertEqual(after["last_change_text"], "v2")
        self.assertIsNotNone(after["last_change_at"])
        self.assertEqual(
            [p["status"] for p in list_polls(self.conn, watch["id"])],
            ["changed", "unchanged", "baseline"])

    def test_pause_preserves_cursor_and_resume_reports_pause_period_change(self):
        fetcher = _Fetcher("v1")
        watch = self._watch()
        run_due_resident_watch(self.conn, watch["id"],
                               fetcher=fetcher, now=self.now)
        cursor = get_watch(self.conn, watch["id"], now=self.now)
        self.assertIsNotNone(cursor["last_content_hash"])
        ok, paused = pause_watch(self.conn, watch["id"], now=self.now)
        self.assertTrue(ok)
        self.assertEqual(paused["status"], "paused")
        self.assertEqual(paused["last_content_hash"],
                         cursor["last_content_hash"])
        self.assertIsNone(paused["next_poll_at"])
        self.assertFalse(paused["due"])
        # 暂停期间即使时间过了间隔，驱动也不轮询。
        self.assertIsNone(run_due_resident_watch(
            self.conn, watch["id"], fetcher=fetcher, now=self.now + 3600))
        self.assertEqual(len(fetcher.calls), 1)
        # 暂停期间网站变了；恢复后第一次轮询用旧游标比对，报出来。
        fetcher.value = "v2"
        ok, resumed = resume_watch(self.conn, watch["id"], now=self.now)
        self.assertTrue(ok)
        self.assertEqual(resumed["status"], "running")
        self.assertEqual(resumed["last_content_hash"],
                         cursor["last_content_hash"])
        poll = run_due_resident_watch(self.conn, watch["id"],
                                      fetcher=fetcher, now=self.now)
        self.assertEqual(poll["status"], "changed")
        self.assertEqual(len(self._tasks()), 1)

    def test_pause_resume_on_wrong_status_are_rejected(self):
        watch = self._watch()
        ok, _ = resume_watch(self.conn, watch["id"], now=self.now)
        self.assertFalse(ok)
        ok, _ = pause_watch(self.conn, watch["id"], now=self.now)
        self.assertTrue(ok)
        ok, _ = pause_watch(self.conn, watch["id"], now=self.now)
        self.assertFalse(ok)

    def test_end_is_terminal_and_name_reusable(self):
        watch = self._watch()
        ok, ended = end_watch(self.conn, watch["id"], now=self.now)
        self.assertTrue(ok)
        self.assertEqual(ended["status"], "ended")
        self.assertIsNone(ended["next_poll_at"])
        for fn in (pause_watch, resume_watch, end_watch):
            ok, old = fn(self.conn, watch["id"], now=self.now)
            self.assertFalse(ok)
            self.assertEqual(old["status"], "ended")
        self.assertIsNone(run_due_resident_watch(
            self.conn, watch["id"], fetcher=_Fetcher(), now=self.now + 3600))
        ok, _ = update_watch(self.conn, watch["id"],
                             {"prompt": "新提示"}, now=self.now)
        self.assertFalse(ok)
        # 已结束的名字可重用。
        again = self._watch()
        self.assertEqual(again["status"], "running")

    def test_claim_is_single_winner(self):
        watch = self._watch()
        self.assertTrue(claim_due_watch(self.conn, watch["id"], now=self.now))
        self.assertFalse(claim_due_watch(self.conn, watch["id"],
                                         now=self.now))
        # 下个间隔到了才能再抢（claim 已把 next_poll_at 预置到下间隔）。
        self.assertTrue(claim_due_watch(self.conn, watch["id"],
                                        now=self.now + 61))

    def test_fetch_failure_backoff_and_cursor_unchanged(self):
        watch = self._watch()
        run_due_resident_watch(self.conn, watch["id"],
                               fetcher=_Fetcher("v1"), now=self.now)
        cursor = get_watch(self.conn, watch["id"], now=self.now)
        self.now += 61
        poll = run_due_resident_watch(self.conn, watch["id"],
                                      fetcher=_failing_fetcher, now=self.now)
        self.assertEqual(poll["status"], "failed")
        after = get_watch(self.conn, watch["id"], now=self.now)
        self.assertEqual(after["consecutive_failures"], 1)
        self.assertEqual(after["last_content_hash"],
                         cursor["last_content_hash"])
        # 失败退避：下次 = now + 间隔 + 30 秒，不是立刻重试。
        self.assertEqual(after["next_poll_at"], self.now + 60 + 30)
        self.now = after["next_poll_at"]
        run_due_resident_watch(self.conn, watch["id"],
                               fetcher=_failing_fetcher, now=self.now)
        after2 = get_watch(self.conn, watch["id"], now=self.now)
        self.assertEqual(after2["consecutive_failures"], 2)
        self.assertEqual(after2["next_poll_at"], self.now + 60 + 60)

    def test_quota_exceeded_pauses_not_ends_and_keeps_cursor(self):
        fetcher = _Fetcher("v1")
        watch = self._watch()
        run_due_resident_watch(self.conn, watch["id"],
                               fetcher=fetcher, now=self.now)
        cursor_hash = get_watch(self.conn, watch["id"],
                                now=self.now)["last_content_hash"]
        self._set_quota("owner1", 0)
        fetcher.value = "v2"
        self.now += 61
        poll = run_due_resident_watch(self.conn, watch["id"],
                                      fetcher=fetcher, now=self.now)
        self.assertEqual(poll["status"], "changed_quota_paused")
        after = get_watch(self.conn, watch["id"], now=self.now)
        self.assertEqual(after["status"], "paused")
        self.assertEqual(after["paused_reason"], "quota_exceeded")
        self.assertNotEqual(after["status"], "ended")
        self.assertEqual(after["last_content_hash"], cursor_hash)
        self.assertEqual(len(self._tasks()), 0)
        # 恢复后配额仍耗尽：再次转暂停，不擅自结束、不丢变化。
        resume_watch(self.conn, watch["id"], now=self.now)
        poll = run_due_resident_watch(self.conn, watch["id"],
                                      fetcher=fetcher, now=self.now)
        self.assertEqual(poll["status"], "changed_quota_paused")
        # 配额恢复后，同一变化仍能报出来（游标没被配额暂停推进）。
        self._set_quota("owner1", 100000000)
        resume_watch(self.conn, watch["id"], now=self.now)
        poll = run_due_resident_watch(self.conn, watch["id"],
                                      fetcher=fetcher, now=self.now)
        self.assertEqual(poll["status"], "changed")
        self.assertEqual(len(self._tasks()), 1)

    def test_update_watch(self):
        watch = self._watch()
        ok, updated = update_watch(
            self.conn, watch["id"],
            {"interval_seconds": 3600, "prompt": "新处理提示"},
            now=self.now)
        self.assertTrue(ok)
        self.assertEqual(updated["interval_seconds"], 3600)
        self.assertEqual(updated["prompt"], "新处理提示")
        with self.assertRaises(ValueError):
            update_watch(self.conn, watch["id"], {}, now=self.now)
        with self.assertRaises(ValueError):
            update_watch(self.conn, watch["id"],
                         {"target_url": "https://localhost/x"}, now=self.now)

    def test_process_due_and_worker_driver(self):
        fetcher = _Fetcher("v1")
        self._watch(name="a")
        self._watch(name="b")
        polls = process_due_resident_watches(self.conn, fetcher=fetcher,
                                             now=self.now)
        self.assertEqual(len(polls), 2)
        self.assertEqual([p["status"] for p in polls],
                         ["baseline", "baseline"])
        self.assertEqual(process_due_resident_watches(
            self.conn, fetcher=fetcher, now=self.now), [])
        # worker 驱动入口：注入假抓取，变化时产出 changed 轮询。
        worker = Worker(self.db_path, "/fake/host.py",
                        env={"XCM_API_KEY": "test-placeholder"},
                        runner=lambda cmd: SimpleNamespace(
                            returncode=0, stdout="{}", stderr=""))
        fetcher.value = "v2"
        self.now += 61
        polls = worker.process_due_resident_watches(
            now=self.now, fetcher=fetcher)
        self.assertEqual(len(polls), 2)
        self.assertTrue(all(p["status"] == "changed" for p in polls))
        self.assertEqual(len(self._tasks()), 2)

    def test_url_validation_blocks_localhost_and_numeric_forms(self):
        for bad in ("http://localhost./x", "http://foo.localhost/x",
                    "http://2130706433/x", "http://0x7f000001/x",
                    "http://0177.0.0.1/x", "http://0x7f.0.0.1/x"):
            with self.assertRaises(ValueError, msg=bad):
                validate_target_url(bad)
        # 规范点分十进制公网 IP 与带尾点的普通域名仍放行。
        self.assertEqual(validate_target_url("https://8.8.8.8/x"),
                         "https://8.8.8.8/x")
        self.assertEqual(validate_target_url("https://example.com./x"),
                         "https://example.com./x")

    def test_unique_index_guards_concurrent_name_clash(self):
        idx = self.conn.execute(
            "SELECT name FROM sqlite_master WHERE type = 'index'"
            " AND name = 'idx_resident_watches_name'").fetchone()
        self.assertIsNotNone(idx)
        self._watch()
        # 正常路径仍先报 ValueError。
        with self.assertRaises(ValueError):
            self._watch()
        # 绕过 SELECT 查重直接插：唯一索引本身必须拦住，
        # 这才是并发下真正的关口。
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO resident_watches(id, project_id, name,"
                " target_url, prompt, interval_seconds, status,"
                " created_by, created_at, updated_at)"
                " VALUES('dup', 'proj', 'watch1', ?, 'p', 60,"
                " 'running', 'owner1', ?, ?)",
                (URL, self.now, self.now))
        other = self._watch(name="watch2")
        with self.assertRaises(ValueError):
            update_watch(self.conn, other["id"], {"name": "watch1"},
                         now=self.now)

    def test_quota_check_failed_is_transient_not_paused(self):
        fetcher = _Fetcher("v1")
        watch = self._watch()
        run_due_resident_watch(self.conn, watch["id"],
                               fetcher=fetcher, now=self.now)
        cursor_hash = get_watch(self.conn, watch["id"],
                                now=self.now)["last_content_hash"]
        original = resident_module._automation_quota_exceeded
        resident_module._automation_quota_exceeded = lambda conn, owner: (
            "quota_check_failed: boom")
        try:
            fetcher.value = "v2"
            self.now += 61
            poll = run_due_resident_watch(
                self.conn, watch["id"], fetcher=fetcher, now=self.now)
        finally:
            resident_module._automation_quota_exceeded = original
        self.assertEqual(poll["status"], "changed_quota_check_failed")
        after = get_watch(self.conn, watch["id"], now=self.now)
        self.assertEqual(after["status"], "running")
        self.assertEqual(after["last_content_hash"], cursor_hash)
        self.assertEqual(after["consecutive_failures"], 1)
        self.assertEqual(after["next_poll_at"], self.now + 60 + 30)
        self.assertEqual(len(self._tasks()), 0)

    def test_pause_during_fetch_aborts_without_task(self):
        fetcher = _Fetcher("v1")
        watch = self._watch()
        run_due_resident_watch(self.conn, watch["id"],
                               fetcher=fetcher, now=self.now)
        cursor_hash = get_watch(self.conn, watch["id"],
                                now=self.now)["last_content_hash"]
        db_path = self.db_path

        def pausing_fetcher(url, timeout):
            conn2 = sqlite3.connect(db_path)
            try:
                conn2.row_factory = sqlite3.Row
                pause_watch(conn2, watch["id"], now=self.now + 61)
            finally:
                conn2.close()
            return "v2"

        self.now += 61
        poll = run_due_resident_watch(
            self.conn, watch["id"], fetcher=pausing_fetcher, now=self.now)
        self.assertIsNone(poll)
        after = get_watch(self.conn, watch["id"], now=self.now)
        self.assertEqual(after["status"], "paused")
        self.assertIsNone(after["next_poll_at"])
        self.assertEqual(after["last_content_hash"], cursor_hash)
        self.assertEqual(len(self._tasks()), 0)
        self.assertEqual(
            [p["status"] for p in list_polls(self.conn, watch["id"])],
            ["baseline"])


class RedirectFetchTestCase(unittest.TestCase):
    """重定向与 DNS 门校验：默认解析器下 redirect-test.test
    经 getaddrinfo 补丁指到本地（127.0.0.1），DNS 门必须拦
    住它（私网解析结果）；相对跳转的跟随逻辑走注入的测试
    resolver（返回公网测试 IP，门正常放行）+ socket 层重定向
    （测试进程内把该公网 IP:端口指回本地服务——门与 pin
    逻辑全程真实跑，只有传输层终点是测试替身）；no_proxy
    同款处理，tearDown 全部还原。"""

    HOST = "redirect-test.test"

    def setUp(self):
        class Handler(http.server.BaseHTTPRequestHandler):
            def do_GET(self):
                if self.path == "/go":
                    self.send_response(302)
                    self.send_header("Location", "/final")
                elif self.path == "/bad":
                    self.send_response(302)
                    self.send_header(
                        "Location", "http://127.0.0.1:1/secret")
                else:
                    body = b"final-body"
                    self.send_response(200)
                    self.send_header("Content-Length", str(len(body)))
                    self.end_headers()
                    self.wfile.write(body)
                    return
                self.end_headers()

            def log_message(self, *args):
                pass

        self.server = http.server.ThreadingHTTPServer(
            ("127.0.0.1", 0), Handler)
        self.port = self.server.server_address[1]
        self.thread = threading.Thread(
            target=self.server.serve_forever, daemon=True)
        self.thread.start()
        self._orig_getaddrinfo = socket.getaddrinfo

        def patched(host, port, family=0, type=0, proto=0, flags=0):
            if host == self.HOST:
                return [(socket.AF_INET, socket.SOCK_STREAM, 6, "",
                         ("127.0.0.1", port))]
            return self._orig_getaddrinfo(
                host, port, family, type, proto, flags)

        socket.getaddrinfo = patched
        # socket 层测试替身：只重定向公网测试 IP 到本地服务端口，
        # 其它地址一律走原始 create_connection。
        self._orig_create_connection = socket.create_connection
        test_port = self.port

        def patched_create_connection(address, *args, **kwargs):
            if address[0] == "93.184.216.34" \
                    and address[1] == test_port:
                address = ("127.0.0.1", test_port)
            return self._orig_create_connection(address, *args, **kwargs)

        socket.create_connection = patched_create_connection
        self._old_no_proxy = os.environ.get("no_proxy")
        self._old_NO_PROXY = os.environ.get("NO_PROXY")
        os.environ["no_proxy"] = (
            (self._old_no_proxy + ",") if self._old_no_proxy
            else "") + self.HOST
        os.environ["NO_PROXY"] = (
            (self._old_NO_PROXY + ",") if self._old_NO_PROXY
            else "") + self.HOST

    def tearDown(self):
        socket.create_connection = self._orig_create_connection
        socket.getaddrinfo = self._orig_getaddrinfo
        if self._old_no_proxy is None:
            os.environ.pop("no_proxy", None)
        else:
            os.environ["no_proxy"] = self._old_no_proxy
        if self._old_NO_PROXY is None:
            os.environ.pop("NO_PROXY", None)
        else:
            os.environ["NO_PROXY"] = self._old_NO_PROXY
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=5)

    def test_relative_redirect_followed_and_private_blocked(self):
        base = "http://%s:%d" % (self.HOST, self.port)
        # 注入 resolver 的跟随路径：解析门放行公网测试 IP，
        # socket 替身把建连指回本地，相对跳转 /go -> /final 走通。
        def test_resolver(host, port):
            return ["93.184.216.34"]

        self.assertEqual(
            default_fetch(base + "/go", 10, resolver=test_resolver),
            "final-body")
        # 默认 DNS 门：测试主机名解析到 127.0.0.1，直接拦住。
        with self.assertRaises(SSRFBlockedError):
            default_fetch(base + "/go", 10)
        # 重定向跳本身：跳转目标是字面私网 IP，逐跳校验拦住。
        with self.assertRaises(ValueError):
            default_fetch(base + "/bad", 10, resolver=test_resolver)
        # 混合解析结果：任一地址私网即整批拒绝，不取公网那个蒙混。
        def mixed(host, port):
            return ["93.184.216.34", "127.0.0.1"]

        with self.assertRaises(ValueError):
            default_fetch(base + "/go", 10, resolver=mixed)


if __name__ == "__main__":
    unittest.main()