devplacepy/tests/unit/services/jobs/deepsearch/enhance.py

# retoor <retoor@molodetz.nl>

import json

from tests.conftest import run_async

from devplacepy.services.jobs.deepsearch import enhance as enhance_module
from devplacepy.services.jobs.deepsearch.enhance import plan_followup_queries


class _FakeResponse:
    def __init__(self, status_code, payload):
        self.status_code = status_code
        self._payload = payload

    def json(self):
        return self._payload


class _FakeClient:
    def __init__(self, response):
        self._response = response

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    async def post(self, *args, **kwargs):
        return self._response


def _client_returning(response):
    def factory(**kwargs):
        return _FakeClient(response)

    return factory


def test_plan_followup_queries_empty_without_covered_titles():
    result = run_async(plan_followup_queries("q", [], "k"))
    assert result == []


def test_plan_followup_queries_parses_gateway_response(monkeypatch):
    payload = {"choices": [{"message": {"content": json.dumps({"queries": ["a", "b"]})}}]}
    monkeypatch.setattr(
        enhance_module.stealth,
        "stealth_async_client",
        _client_returning(_FakeResponse(200, payload)),
    )
    result = run_async(plan_followup_queries("q", ["Existing source"], "k"))
    assert result == ["a", "b"]


def test_plan_followup_queries_empty_when_sources_already_cover_question(monkeypatch):
    payload = {"choices": [{"message": {"content": json.dumps({"queries": []})}}]}
    monkeypatch.setattr(
        enhance_module.stealth,
        "stealth_async_client",
        _client_returning(_FakeResponse(200, payload)),
    )
    result = run_async(plan_followup_queries("q", ["Existing source"], "k"))
    assert result == []


def test_plan_followup_queries_fails_soft_on_gateway_error(monkeypatch):
    monkeypatch.setattr(
        enhance_module.stealth,
        "stealth_async_client",
        _client_returning(_FakeResponse(500, {})),
    )
    result = run_async(plan_followup_queries("q", ["Existing source"], "k"))
    assert result == []


def test_plan_followup_queries_caps_at_max(monkeypatch):
    payload = {
        "choices": [
            {"message": {"content": json.dumps({"queries": ["a", "b", "c", "d", "e"]})}}
        ]
    }
    monkeypatch.setattr(
        enhance_module.stealth,
        "stealth_async_client",
        _client_returning(_FakeResponse(200, payload)),
    )
    result = run_async(plan_followup_queries("q", ["Existing source"], "k"))
    assert len(result) == enhance_module.MAX_FOLLOWUP_QUERIES