tests/test_image_inversion_job_api.py

from __future__ import annotations

import ast
import io
import logging
import os
import socket
import tempfile
import threading
import time
import unittest
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from unittest.mock import patch

import httpx
from fastapi import FastAPI
from fastapi.testclient import TestClient
from PIL import Image
from pydantic import ValidationError

from image_inversion_job_api import (
    ImageInversionJobAPI,
    ImageInversionJobRecord,
    ImageInversionJobRequest,
    ImageInversionJobSettings,
    ImageSourceError,
    create_job_api_lifespan,
    decode_image_bytes,
    fetch_remote_image,
    fingerprint_request,
    validate_image_url,
)


API_KEY = "k" * 48
EXPECTED_RESULT = {
    "status_markdown": "done",
    "general_tags_html": "<p>general</p>",
    "character_tags_html": "<p>character</p>",
    "ip_tags_html": "<p>ip</p>",
    "summary_text": "tag one, tag two",
    "metadata": {"device": "cpu", "latency_s_total": 0.1},
}


def _make_png_bytes(size: tuple[int, int] = (96, 80)) -> bytes:
    """生成远程抓图测试使用的有效 PNG。

    Args:
        size: PNG 的宽高。

    Returns:
        完整 PNG 文件字节。
    """
    image_buffer = io.BytesIO()
    with Image.new("RGB", size, "green") as image:
        image.save(image_buffer, format="PNG")
    return image_buffer.getvalue()


PNG_BYTES = _make_png_bytes()


class _TrackedImage:
    """记录任务输入图是否被 API 生命周期可靠关闭。"""

    def __init__(self) -> None:
        self.closed = False

    def close(self) -> None:
        """标记当前测试图片已经关闭。"""
        self.closed = True


class _FakeStreamResponse:
    """提供 httpx 流式响应所需的最小测试替身。"""

    def __init__(
        self,
        status_code: int,
        body: bytes,
        *,
        peer_address: str = "93.184.216.34",
        extra_headers: dict[str, str] | None = None,
    ) -> None:
        self.status_code = status_code
        self.body = body
        self.headers = {
            "content-type": "image/png",
            "content-length": str(len(body)),
        }
        if extra_headers:
            self.headers.update(extra_headers)
        self.extensions = {
            "network_stream": type(
                "NetworkStreamStub",
                (),
                {
                    "get_extra_info": lambda self, name: (
                        (peer_address, 443) if name == "server_addr" else None
                    )
                },
            )()
        }

    def __enter__(self) -> "_FakeStreamResponse":
        """进入流式响应上下文并返回自身。"""
        return self

    def __exit__(self, exc_type, exc, traceback) -> bool:
        """退出响应上下文且不吞掉异常。"""
        del exc_type, exc, traceback
        return False

    def iter_bytes(self):
        """以单个分块返回完整响应体。"""
        yield self.body


class ImageInversionJobAPITest(unittest.TestCase):
    """验证 Image Inversion 自定义任务 API 的公开契约与安全边界。"""

    def setUp(self) -> None:
        """登记每个测试需要在结束时关闭的客户端和线程事件。"""
        self.clients: list[TestClient] = []
        self.release_events: list[threading.Event] = []

    def tearDown(self) -> None:
        """释放测试线程和 TestClient，不创建持久化测试产物。"""
        for release_event in self.release_events:
            release_event.set()
        time.sleep(0.02)
        for client in self.clients:
            client.close()

    @staticmethod
    def _settings(**overrides) -> ImageInversionJobSettings:
        """构造不依赖进程环境的隔离任务设置。

        Args:
            overrides: 要覆盖的 ImageInversionJobSettings 字段。

        Returns:
            可供单元测试直接注入的设置。
        """
        values = {
            "api_key": API_KEY,
            "allowed_hosts": frozenset({"allowed.example"}),
            "result_ttl_seconds": 1800,
            "poll_after_seconds": 1,
            "fetch_timeout_seconds": 15.0,
            "max_image_bytes": 8 * 1024 * 1024,
            "max_image_pixels": 4_194_304,
            "max_image_dimension": 4096,
            "max_records": 256,
            "space_host": "https://space.example",
        }
        values.update(overrides)
        return ImageInversionJobSettings(**values)

    @staticmethod
    def _payload(**overrides) -> dict:
        """生成可直接提交给创建接口的合法请求体。

        Args:
            overrides: 要覆盖或增加的 JSON 字段。

        Returns:
            包含远程输入图和全部展示参数的请求字典。
        """
        payload = {
            "input_image_url": "https://allowed.example/input.png?source=private",
            "general_threshold": 0.35,
            "character_threshold": 0.85,
            "show_confidence": True,
            "show_general": True,
            "show_character": True,
            "show_ip": True,
            "separator": "comma",
            "show_chinese": True,
        }
        payload.update(overrides)
        return payload

    @staticmethod
    def _headers(**overrides) -> dict[str, str]:
        """生成包含共享密钥的请求头。

        Args:
            overrides: 要覆盖或追加的请求头。

        Returns:
            可传给 TestClient 的请求头字典。
        """
        headers = {"X-API-Key": API_KEY}
        headers.update(overrides)
        return headers

    def _make_service(
        self,
        *,
        executor=None,
        image_fetcher=None,
        inference_slot: threading.Lock | None = None,
        settings: ImageInversionJobSettings | None = None,
        add_catch_all: bool = False,
    ) -> tuple[TestClient, ImageInversionJobAPI, list[_TrackedImage], dict[str, int]]:
        """创建完全隔离且不导入模型的 FastAPI 测试服务。

        Args:
            executor: 可选假标签执行器。
            image_fetcher: 可选假远程图片抓取器。
            inference_slot: 可选 UI/API 共享锁。
            settings: 可选任务设置。
            add_catch_all: 是否先注册模拟 Gradio 的宽泛路由。

        Returns:
            TestClient、API 对象、输入图片列表与调用计数器。
        """
        settings = settings or self._settings()
        created_images: list[_TrackedImage] = []
        counters = {"fetch": 0, "execute": 0}
        counter_lock = threading.Lock()

        if image_fetcher is None:

            def image_fetcher(url, current_settings):
                """返回可追踪关闭状态的假输入图。"""
                del url, current_settings
                with counter_lock:
                    counters["fetch"] += 1
                image = _TrackedImage()
                created_images.append(image)
                return image

        if executor is None:

            def executor(payload, input_image, job_id):
                """返回与新接口契约一致的六字段 JSON 结果。"""
                del payload, input_image, job_id
                with counter_lock:
                    counters["execute"] += 1
                return EXPECTED_RESULT

        api = ImageInversionJobAPI(
            settings=settings,
            executor=executor,
            inference_slot=inference_slot or threading.Lock(),
            image_fetcher=image_fetcher,
        )
        app = FastAPI()
        if add_catch_all:

            @app.api_route("/{path:path}", methods=["GET", "POST"])
            def catch_all(path: str):
                """模拟可能抢先匹配 API 请求的 Gradio 宽泛路由。"""
                return {"catch_all": path}

        api.install_on_app(app)
        client = TestClient(app, base_url="https://space.example")
        self.clients.append(client)
        return client, api, created_images, counters

    def _wait_for_terminal(
        self,
        client: TestClient,
        job_id: str,
        timeout_seconds: float = 3.0,
    ):
        """轮询测试任务直到返回成功或失败终态。

        Args:
            client: 当前隔离服务客户端。
            job_id: 创建接口返回的任务标识。
            timeout_seconds: 允许后台测试线程运行的最长时间。

        Returns:
            第一个非 202 状态响应。
        """
        deadline = time.monotonic() + timeout_seconds
        while time.monotonic() < deadline:
            response = client.get(
                f"/api/jobs/{job_id}",
                headers={"X-API-Key": API_KEY},
            )
            if response.status_code != 202:
                return response
            time.sleep(0.01)
        self.fail(f"job {job_id} did not reach a terminal state")

    def test_auth_schema_cpu_contract_and_route_precedence(self) -> None:
        """验证鉴权、CPU 请求契约及自定义路由优先级。"""
        client, api, _, counters = self._make_service(add_catch_all=True)
        self.assertTrue(callable(create_job_api_lifespan(api)))

        unauthorized = client.post("/api/jobs", json=self._payload())
        self.assertEqual(unauthorized.status_code, 401)
        self.assertEqual(counters["fetch"], 0)

        unknown_field = client.post(
            "/api/jobs",
            json=self._payload(unknown=True),
            headers=self._headers(),
        )
        self.assertEqual(unknown_field.status_code, 422)
        self.assertEqual(counters["fetch"], 0)

        # cpu-basic Space 不依赖 ZeroGPU 的 x-ip-token。
        accepted = client.post(
            "/api/jobs",
            json=self._payload(),
            headers=self._headers(),
        )
        self.assertEqual(accepted.status_code, 202)
        self.assertNotIn("catch_all", accepted.json())
        terminal = self._wait_for_terminal(client, accepted.json()["job_id"])
        self.assertEqual(terminal.status_code, 200)
        self.assertEqual(terminal.json(), EXPECTED_RESULT)

    def test_success_contract_does_not_retain_caller_credentials(self) -> None:
        """验证成功链路仅向执行器传业务参数且结果无包装层。"""
        seen: dict = {}

        def executor(payload, input_image, job_id):
            """记录后台适配器参数并返回固定六字段结果。"""
            seen.update({"payload": payload, "input_image": input_image, "job_id": job_id})
            return EXPECTED_RESULT

        client, api, images, _ = self._make_service(executor=executor)
        created = client.post(
            "/api/jobs",
            json=self._payload(separator="newline"),
            headers=self._headers(
                Authorization="Bearer must-not-be-retained",
                Cookie="must-not-be-retained",
                **{"X-IP-Token": "not-needed-on-cpu"},
            ),
        )
        self.assertEqual(created.status_code, 202)
        self.assertEqual(created.headers["Cache-Control"], "no-store")
        self.assertEqual(created.headers["Location"], created.json()["status_url"])
        job_id = created.json()["job_id"]
        terminal = self._wait_for_terminal(client, job_id)
        self.assertEqual(terminal.status_code, 200)
        self.assertEqual(set(terminal.json()), set(EXPECTED_RESULT))
        self.assertEqual(seen["job_id"], job_id)
        self.assertEqual(seen["payload"].separator, "newline")
        self.assertTrue(all(image.closed for image in images))
        self.assertFalse(hasattr(api.jobs[job_id], "zero_gpu_headers"))

    def test_idempotent_replay_conflict_and_busy_submission(self) -> None:
        """验证幂等重放、冲突及单任务繁忙状态。"""
        release_event = threading.Event()
        started_event = threading.Event()
        self.release_events.append(release_event)
        execution_count = 0
        count_lock = threading.Lock()

        def blocking_executor(payload, input_image, job_id):
            """阻塞首个任务，使测试可观察幂等和忙状态。"""
            nonlocal execution_count
            del payload, input_image, job_id
            with count_lock:
                execution_count += 1
            started_event.set()
            release_event.wait(timeout=3)
            return EXPECTED_RESULT

        client, _, _, counters = self._make_service(executor=blocking_executor)
        headers = self._headers(**{"Idempotency-Key": "same-request"})
        first = client.post("/api/jobs", json=self._payload(), headers=headers)
        self.assertEqual(first.status_code, 202)
        self.assertTrue(started_event.wait(timeout=1))

        replay = client.post("/api/jobs", json=self._payload(), headers=headers)
        self.assertEqual(replay.status_code, 202)
        self.assertEqual(replay.json()["job_id"], first.json()["job_id"])
        self.assertEqual(counters["fetch"], 1)

        conflict = client.post(
            "/api/jobs",
            json=self._payload(general_threshold=0.5),
            headers=headers,
        )
        self.assertEqual(conflict.status_code, 409)

        busy = client.post(
            "/api/jobs",
            json=self._payload(character_threshold=0.7),
            headers=self._headers(**{"Idempotency-Key": "different-request"}),
        )
        self.assertEqual(busy.status_code, 503)
        self.assertEqual(busy.headers["Retry-After"], "5")
        self.assertEqual(counters["fetch"], 1)

        release_event.set()
        terminal = self._wait_for_terminal(client, first.json()["job_id"])
        self.assertEqual(terminal.status_code, 200)
        self.assertEqual(execution_count, 1)

    def test_concurrent_same_idempotency_key_executes_once(self) -> None:
        """验证并发同键请求最终只执行一次。"""
        fetch_barrier = threading.Barrier(2)
        fetch_count = 0
        execute_count = 0
        count_lock = threading.Lock()

        def racing_fetcher(url, settings):
            """让两个请求都完成第一次幂等检查后再返回图片。"""
            nonlocal fetch_count
            del url, settings
            with count_lock:
                fetch_count += 1
            fetch_barrier.wait(timeout=2)
            return _TrackedImage()

        def counting_executor(payload, input_image, job_id):
            """统计实际执行次数并返回固定结果。"""
            nonlocal execute_count
            del payload, input_image, job_id
            with count_lock:
                execute_count += 1
            return EXPECTED_RESULT

        client, _, _, _ = self._make_service(
            image_fetcher=racing_fetcher,
            executor=counting_executor,
        )
        headers = self._headers(**{"Idempotency-Key": "concurrent-key"})

        def submit_job(_index: int):
            """提交一个并发幂等请求。"""
            return client.post("/api/jobs", json=self._payload(), headers=headers)

        with ThreadPoolExecutor(max_workers=2) as pool:
            responses = list(pool.map(submit_job, range(2)))
        self.assertEqual([response.status_code for response in responses], [202, 202])
        self.assertEqual(responses[0].json()["job_id"], responses[1].json()["job_id"])
        terminal = self._wait_for_terminal(client, responses[0].json()["job_id"])
        self.assertEqual(terminal.status_code, 200)
        self.assertEqual(fetch_count, 2)
        self.assertEqual(execute_count, 1)

    def test_failure_is_sanitized_and_releases_image_and_slot(self) -> None:
        """验证内部异常不泄密，且图片与共享锁始终释放。"""
        shared_slot = threading.Lock()

        def failing_executor(payload, input_image, job_id):
            """模拟包含敏感内部信息的预测错误。"""
            del payload, input_image, job_id
            raise RuntimeError("INTERNAL-SECRET-STACK-DATA")

        client, api, images, _ = self._make_service(
            executor=failing_executor,
            inference_slot=shared_slot,
        )
        with self.assertLogs("image_inversion_job_api", level=logging.ERROR) as logs:
            created = client.post(
                "/api/jobs",
                json=self._payload(),
                headers=self._headers(Authorization="Bearer sensitive"),
            )
            terminal = self._wait_for_terminal(client, created.json()["job_id"])

        self.assertEqual(terminal.status_code, 500)
        self.assertEqual(terminal.json(), {"error": {"code": "PREDICTION_FAILED"}})
        combined_logs = "\n".join(logs.output)
        self.assertNotIn("INTERNAL-SECRET-STACK-DATA", combined_logs)
        self.assertNotIn(API_KEY, combined_logs)
        self.assertTrue(all(image.closed for image in images))
        self.assertTrue(shared_slot.acquire(blocking=False))
        shared_slot.release()
        self.assertEqual(api.jobs[created.json()["job_id"]].error_type, "RuntimeError")

    def test_thread_start_failure_rolls_back_and_retry_succeeds(self) -> None:
        """验证后台线程启动失败时任务、幂等索引和锁均回滚。"""
        client, api, images, counters = self._make_service()
        with patch(
            "image_inversion_job_api.threading.Thread.start",
            side_effect=RuntimeError("thread-start-failed"),
        ):
            with self.assertRaises(RuntimeError):
                client.post(
                    "/api/jobs",
                    json=self._payload(),
                    headers=self._headers(**{"Idempotency-Key": "rollback-key"}),
                )

        self.assertEqual(api.jobs, {})
        self.assertEqual(api.idempotency_jobs, {})
        self.assertTrue(all(image.closed for image in images))
        self.assertFalse(api.inference_slot.locked())

        retry = client.post(
            "/api/jobs",
            json=self._payload(),
            headers=self._headers(**{"Idempotency-Key": "rollback-key"}),
        )
        self.assertEqual(retry.status_code, 202)
        self.assertEqual(
            self._wait_for_terminal(client, retry.json()["job_id"]).status_code,
            200,
        )
        self.assertEqual(counters["execute"], 1)

    def test_ttl_removes_all_expired_terminal_records(self) -> None:
        """验证 TTL 清理会移除全部已过期终态记录。"""
        settings = self._settings(result_ttl_seconds=60)
        client, api, _, _ = self._make_service(settings=settings)
        job_ids: list[str] = []
        for threshold in (0.2, 0.3):
            created = client.post(
                "/api/jobs",
                json=self._payload(general_threshold=threshold),
                headers=self._headers(),
            )
            job_id = created.json()["job_id"]
            self.assertEqual(self._wait_for_terminal(client, job_id).status_code, 200)
            api.jobs[job_id].completed_at = time.time() - 100
            job_ids.append(job_id)

        api.remove_expired_jobs()
        self.assertEqual(api.jobs, {})
        self.assertEqual(len(job_ids), 2)

    def test_capacity_evicts_oldest_terminal_and_preserves_running_records(self) -> None:
        """验证容量上限只淘汰最老终态且满是运行态时返回 503。"""
        settings = self._settings(max_records=2)
        client, api, _, _ = self._make_service(settings=settings)
        completed_job_ids: list[str] = []
        for threshold in (0.2, 0.3):
            created = client.post(
                "/api/jobs",
                json=self._payload(general_threshold=threshold),
                headers=self._headers(),
            )
            job_id = created.json()["job_id"]
            self.assertEqual(self._wait_for_terminal(client, job_id).status_code, 200)
            completed_job_ids.append(job_id)
            time.sleep(0.01)

        third = client.post(
            "/api/jobs",
            json=self._payload(general_threshold=0.4),
            headers=self._headers(),
        )
        self.assertEqual(third.status_code, 202)
        self.assertNotIn(completed_job_ids[0], api.jobs)
        self.assertIn(completed_job_ids[1], api.jobs)
        self.assertLessEqual(len(api.jobs), 2)
        self.assertEqual(self._wait_for_terminal(client, third.json()["job_id"]).status_code, 200)

        full_settings = self._settings(max_records=1)
        full_client, full_api, _, counters = self._make_service(settings=full_settings)
        running_record = ImageInversionJobRecord(
            job_id="running-record",
            payload=ImageInversionJobRequest(**self._payload()),
            request_fingerprint="fingerprint",
            input_image=_TrackedImage(),
            status="running",
        )
        full_api.jobs[running_record.job_id] = running_record
        blocked = full_client.post(
            "/api/jobs",
            json=self._payload(general_threshold=0.6),
            headers=self._headers(),
        )
        self.assertEqual(blocked.status_code, 503)
        self.assertEqual(blocked.headers["Retry-After"], "5")
        self.assertIn(running_record.job_id, full_api.jobs)
        self.assertEqual(counters["fetch"], 0)

    def test_settings_schema_and_fingerprint(self) -> None:
        """验证环境设置、公开 Schema 和请求指纹边界。"""
        with patch.dict(
            os.environ,
            {
                "JOB_API_KEY": API_KEY,
                "JOB_IMAGE_ALLOWED_HOSTS": "Allowed.Example,cdn.example",
            },
            clear=True,
        ):
            settings = ImageInversionJobSettings.from_env()
        self.assertEqual(
            settings.allowed_hosts,
            frozenset({"allowed.example", "cdn.example"}),
        )
        self.assertEqual(settings.max_image_bytes, 8 * 1024 * 1024)
        self.assertEqual(settings.max_records, 256)

        for invalid_max_records in (7, 4097):
            with self.subTest(max_records=invalid_max_records), patch.dict(
                os.environ,
                {
                    "JOB_API_KEY": API_KEY,
                    "JOB_IMAGE_ALLOWED_HOSTS": "allowed.example",
                    "JOB_MAX_RECORDS": str(invalid_max_records),
                },
                clear=True,
            ):
                with self.assertRaises(RuntimeError):
                    ImageInversionJobSettings.from_env()

        minimum = ImageInversionJobRequest(
            input_image_url="https://allowed.example/input.png",
            general_threshold=0.0,
            character_threshold=0.0,
            separator="space",
        )
        maximum = ImageInversionJobRequest(
            input_image_url="https://allowed.example/input.png",
            general_threshold=1.0,
            character_threshold=1.0,
            separator="newline",
        )
        self.assertNotEqual(fingerprint_request(minimum), fingerprint_request(maximum))
        reordered = ImageInversionJobRequest.model_validate(
            minimum.model_dump(mode="json")
        )
        self.assertEqual(fingerprint_request(minimum), fingerprint_request(reordered))

        for invalid_payload in (
            self._payload(input_image_url="http://allowed.example/input.png"),
            self._payload(general_threshold=-0.01),
            self._payload(character_threshold=1.01),
            self._payload(general_threshold=float("nan")),
            self._payload(separator="逗号"),
            self._payload(unknown=True),
        ):
            with self.subTest(payload=invalid_payload):
                with self.assertRaises(ValidationError):
                    ImageInversionJobRequest(**invalid_payload)

    def test_url_policy_decode_and_fetch_do_not_leak_headers(self) -> None:
        """验证白名单、SSRF、图片解码和无凭据抓取边界。"""
        settings = self._settings()
        public_record = (
            socket.AF_INET,
            socket.SOCK_STREAM,
            socket.IPPROTO_TCP,
            "",
            ("93.184.216.34", 443),
        )
        with patch(
            "image_inversion_job_api.socket.getaddrinfo",
            return_value=[public_record],
        ):
            validated = validate_image_url(
                "https://allowed.example/image.png?private=query",
                settings,
            )
            self.assertEqual(validated.host, "allowed.example")
            for invalid_url in (
                "http://allowed.example/image.png",
                "https://user:password@allowed.example/image.png",
                "https://allowed.example:444/image.png",
                "https://sub.allowed.example/image.png",
                "https://allowed.example/image.png#fragment",
            ):
                with self.subTest(url=invalid_url):
                    with self.assertRaises(ImageSourceError):
                        validate_image_url(invalid_url, settings)

        private_record = (
            socket.AF_INET,
            socket.SOCK_STREAM,
            socket.IPPROTO_TCP,
            "",
            ("127.0.0.1", 443),
        )
        with patch(
            "image_inversion_job_api.socket.getaddrinfo",
            return_value=[public_record, private_record],
        ):
            with self.assertRaises(ImageSourceError):
                validate_image_url("https://allowed.example/image.png", settings)

        decoded = decode_image_bytes(PNG_BYTES, "application/octet-stream", settings)
        self.assertEqual(decoded.mode, "RGB")
        decoded.close()

        captured: dict = {}

        class FakeClient:
            """捕获 httpx 客户端配置并返回假响应。"""

            def __init__(self, **kwargs) -> None:
                captured["client_kwargs"] = kwargs

            def __enter__(self):
                """进入客户端上下文并返回自身。"""
                return self

            def __exit__(self, exc_type, exc, traceback) -> bool:
                """退出客户端上下文且不吞掉异常。"""
                del exc_type, exc, traceback
                return False

            def close(self) -> None:
                """模拟关闭客户端。"""
                captured["closed"] = True

            def stream(self, method, url, headers):
                """记录公开请求参数并返回假响应。"""
                captured["method"] = method
                captured["url"] = url
                captured["headers"] = dict(headers)
                return captured["response"]

        with patch(
            "image_inversion_job_api.socket.getaddrinfo",
            return_value=[public_record],
        ), patch("image_inversion_job_api.httpx.Client", FakeClient):
            captured["response"] = _FakeStreamResponse(200, PNG_BYTES)
            fetched = fetch_remote_image(
                "https://allowed.example/image.png?secret=value",
                settings,
            )
            fetched.close()
            self.assertFalse(captured["client_kwargs"]["follow_redirects"])
            self.assertFalse(captured["client_kwargs"]["trust_env"])
            self.assertEqual(
                captured["headers"],
                {"Accept": "image/png,image/jpeg,image/webp"},
            )
            for header_name in ("Authorization", "Cookie", "X-API-Key", "X-IP-Token"):
                self.assertNotIn(header_name, captured["headers"])

            captured["response"] = _FakeStreamResponse(302, b"")
            with self.assertRaises(ImageSourceError):
                fetch_remote_image("https://allowed.example/redirect", settings)

            captured["response"] = _FakeStreamResponse(
                200,
                PNG_BYTES,
                peer_address="127.0.0.1",
            )
            with self.assertRaises(ImageSourceError):
                fetch_remote_image("https://allowed.example/rebound", settings)

    def test_app_wires_job_lifespan_and_scopes_hf_token(self) -> None:
        """静态验证应用接线、六字段结果与受限模型下载令牌边界。"""
        app_path = Path(__file__).resolve().parents[1] / "app.py"
        app_source = app_path.read_text(encoding="utf-8")
        app_tree = ast.parse(app_source)

        executor_function = next(
            node
            for node in app_tree.body
            if isinstance(node, ast.FunctionDef)
            and node.name == "execute_image_inversion_job"
        )
        result_return = next(
            node
            for node in ast.walk(executor_function)
            if isinstance(node, ast.Return) and isinstance(node.value, ast.Dict)
        )
        result_keys = {
            key.value
            for key in result_return.value.keys
            if isinstance(key, ast.Constant) and isinstance(key.value, str)
        }
        self.assertEqual(result_keys, set(EXPECTED_RESULT))
        self.assertIn("ssr_mode=False", app_source)
        self.assertIn('app_kwargs={"lifespan": JOB_API_LIFESPAN}', app_source)
        self.assertIn('ASSETS_HF_TOKEN = os.environ.get("HF_TOKEN")', app_source)
        self.assertIn('token=ASSETS_HF_TOKEN or False', app_source)
        self.assertIn('HF_HUB_DISABLE_IMPLICIT_TOKEN', app_source)


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