zcbot/tests/test_paper_server_tools.py

import os
import tempfile
import unittest
from contextlib import ExitStack
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
from uuid import uuid4

from platform_sources.paper_server import (
    PaperServerFetchTool,
    PaperServerSearchTool,
    _media_url,
)
from tools.run_python import RunPythonTool


class PaperServerToolTests(unittest.TestCase):
    def test_search_forwards_publication_type_to_server_filter(self):
        captured = {}

        def fake_get_json(url, *, api_key, params=None):
            captured.update(params or {})
            self.assertEqual(api_key, "secret")
            return {"results": [{"id": "p1", "type": "book", "title": "Book"}]}

        with (
            patch(
                "platform_sources.paper_server._config",
                return_value=(
                    "https://paper.test",
                    "https://paper.test/api/resm/paper",
                    "secret",
                ),
            ),
            patch("platform_sources.paper_server._get_json", side_effect=fake_get_json),
        ):
            result = PaperServerSearchTool().execute(publication_type="book", limit=20)

        self.assertEqual(captured["type"], "book")
        self.assertEqual(captured["page_size"], 20)
        self.assertIn('"type": "book"', result)
        self.assertNotIn("secret", result)

    def test_search_accepts_open_raw_type_without_closed_enum(self):
        with (
            patch(
                "platform_sources.paper_server._config",
                return_value=(
                    "https://paper.test",
                    "https://paper.test/api/resm/paper",
                    "secret",
                ),
            ),
            patch("platform_sources.paper_server._get_json", return_value=[]) as request,
        ):
            result = PaperServerSearchTool().execute(publication_type="future-type")
        self.assertEqual(request.call_args.kwargs["params"]["type"], "future-type")
        self.assertEqual(result, "[]")

    def test_search_rejects_unsafe_raw_type(self):
        with patch("platform_sources.paper_server._get_json") as request:
            result = PaperServerSearchTool().execute(publication_type="book&api_key=leak")
        request.assert_not_called()
        self.assertTrue(result.startswith("[Error]"))

    def test_media_download_url_must_match_configured_origin(self):
        self.assertEqual(
            _media_url("/media/a.pdf", "https://paper.test"),
            "https://paper.test/media/a.pdf",
        )
        with self.assertRaisesRegex(RuntimeError, "out-of-origin"):
            _media_url("https://attacker.test/a.pdf", "https://paper.test")

    def test_fetch_schema_limits_format(self):
        with tempfile.TemporaryDirectory() as tmp:
            tool = PaperServerFetchTool(working_dir=Path(tmp))
            self.assertEqual(
                tool.parameters["properties"]["format"]["enum"], ["pdf", "xml"]
            )

    def test_fetch_downloads_atomically_and_reuses_existing_file(self):
        class FakeStreamResponse:
            status_code = 200

            def __enter__(self):
                return self

            def __exit__(self, *_args):
                return False

            def iter_bytes(self, chunk_size=0):
                self.test_chunk_size = chunk_size
                yield b"%PDF-test"

        paper = {
            "id": "p1",
            "doi": "10.1/example",
            "has_fulltext_pdf": True,
            "pdf_url": "/media/example.pdf",
        }
        with (
            tempfile.TemporaryDirectory() as tmp,
            patch("platform_sources.paper_server._get_paper", return_value=paper),
            patch(
                "platform_sources.paper_server._config",
                return_value=(
                    "https://paper.test",
                    "https://paper.test/api/resm/paper",
                    "secret",
                ),
            ),
            patch(
                "platform_sources.paper_server.httpx.stream", return_value=FakeStreamResponse()
            ) as stream,
        ):
            tool = PaperServerFetchTool(working_dir=Path(tmp))
            first = tool.execute(id_or_doi="p1", format="pdf")
            second = tool.execute(id_or_doi="p1", format="pdf")
            destination = Path(tmp) / "papers" / "10.1_example.pdf"
            saved_bytes = destination.read_bytes()

        self.assertEqual(saved_bytes, b"%PDF-test")
        self.assertTrue(first.startswith("saved:"))
        self.assertTrue(second.endswith("(existing)"))
        self.assertEqual(stream.call_count, 1)

    def test_run_python_no_longer_passes_paper_server_key(self):
        with patch.dict(os.environ, {"PAPER_SERVER_API_KEY": "secret"}, clear=False):
            env = RunPythonTool()._filtered_env()
        self.assertNotIn("PAPER_SERVER_API_KEY", env)

    def test_registry_exposes_host_tools_only_when_platform_key_exists(self):
        from core.tool_registry import ToolContext, build_tools

        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            common = dict(
                tool_base=root,
                ur_path=root,
                working_dir_path=root,
                task_id=uuid4(),
                uid=uuid4(),
                cfg={},
                caps=SimpleNamespace(enable_run_python=False),
                skills=SimpleNamespace(skills={}),
                cancel_check=None,
                scheduled_run=False,
                deferred_actions=SimpleNamespace(),
                ark_cfg=None,
                img_provider="",
                img_key="",
                img_cfg=None,
                img_provider_cfg=None,
                video_variant="",
                office_to_pdf_available=False,
            )
            gates = dict(
                DOCUMENT_SEARCH_API_KEY="",
                MP_API_KEY="",
                PAPER_SERVER_API_KEY="",
            )

            def enter_common_patches(stack):
                stack.enter_context(
                    patch(
                        "core.tool_registry._external_system_status_available",
                        return_value=False,
                    )
                )
                stack.enter_context(
                    patch(
                        "core.tool_registry._external_systems_available",
                        return_value=False,
                    )
                )
                stack.enter_context(
                    patch("core.tool_registry.smtp_configured", return_value=False)
                )
                stack.enter_context(
                    patch(
                        "core.tool_registry.wechat_push_available", return_value=False
                    )
                )
                stack.enter_context(
                    patch("core.tool_registry.lfasr_configured", return_value=False)
                )
                stack.enter_context(
                    patch("core.tool_registry.BochaConfig.load", return_value=None)
                )

            with ExitStack() as stack:
                stack.enter_context(patch.dict(os.environ, gates, clear=False))
                enter_common_patches(stack)
                without_key = build_tools(ToolContext(**common))
            gates["PAPER_SERVER_API_KEY"] = "secret"
            with ExitStack() as stack:
                stack.enter_context(patch.dict(os.environ, gates, clear=False))
                enter_common_patches(stack)
                with_key = build_tools(ToolContext(**common))

        self.assertNotIn("paper_server_search", without_key)
        self.assertTrue(
            {"paper_server_search", "paper_server_get", "paper_server_fetch"}.issubset(with_key)
        )


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