TraderAI/tests/test_scheduler.py

import pytest

from traderai.agent import OllamaAgent
from traderai.memory import MemoryStore
from traderai.scheduler import WakeScheduler


class FakeUEXNotifications:
    def __init__(self):
        self.calls = 0

    async def get_user_notifications(self):
        self.calls += 1
        return {
            "status": "ok",
            "notifications": [
                {
                    "id": 10,
                    "message": "A buyer replied to your listing.",
                    "redir": "/marketplace/negotiations/abc",
                    "code": "negotiation_reply",
                    "date_added": 123,
                    "date_read": 0,
                },
                {
                    "id": 11,
                    "message": "Already read.",
                    "date_added": 122,
                    "date_read": 123,
                },
            ],
        }


class FailingUEXNotifications:
    async def get_user_notifications(self):
        raise RuntimeError("bad token")


class FakeWakeAgent:
    async def generate_wake_response(self, wake_message):
        return f"Wake output: {wake_message}"


class ListingWakeTools:
    schemas = [
        {
            "type": "function",
            "function": {
                "name": "search_marketplace_listings",
                "description": "Search active UEX marketplace listings.",
                "parameters": {"type": "object", "properties": {}},
            },
        }
    ]

    def __init__(self):
        self.calls = []
        self.pending_actions = {}

    async def execute(self, name, arguments):
        self.calls.append((name, arguments))
        return {
            "count": 2,
            "listings": [
                {
                    "id": 100,
                    "title": "Wikelo Favor",
                    "operation": "sell",
                    "price": 500_000_000,
                    "currency": "UEC",
                    "in_stock": 9,
                    "advertiser": "pilot_a",
                },
                {
                    "id": 101,
                    "title": "Wikelo Favor stack",
                    "operation": "sell",
                    "price": 1_000_000_000,
                    "currency": "UEC",
                    "in_stock": 5,
                    "advertiser": "pilot_b",
                },
            ],
        }


class ListingWakeAgent(OllamaAgent):
    def __init__(self, memory):
        self.listing_tools = ListingWakeTools()
        super().__init__("http://127.0.0.1:1", "missing-model", self.listing_tools, memory=memory)
        self.responses = [
            {
                "message": {
                    "role": "assistant",
                    "content": "",
                    "tool_calls": [
                        {
                            "function": {
                                "name": "search_marketplace_listings",
                                "arguments": {
                                    "query": "Wikelo Favor",
                                    "operation": "sell",
                                    "limit": 5,
                                },
                            }
                        }
                    ],
                }
            },
            {
                "message": {
                    "role": "assistant",
                    "content": (
                        "Listing check complete: found 2 active Wikelo Favor sell listings. "
                        "Cheapest listing is 500,000,000 UEC with 9 in stock; the next listing is "
                        "1,000,000,000 UEC. Suggested next action: price near 500,000,000 UEC "
                        "if you want to move yours quickly."
                    ),
                }
            },
        ]

    async def ensure_available(self):
        return None

    async def _ollama_chat(self, *args, **kwargs):
        return self.responses.pop(0)


@pytest.mark.asyncio
async def test_poll_uex_notifications_adds_unread_once(tmp_path):
    memory = MemoryStore(str(tmp_path / "memory.sqlite3"))
    scheduler = WakeScheduler(memory)
    scheduler.bind_uex_notifications(FakeUEXNotifications())

    first = await scheduler.poll_uex_notifications()
    second = await scheduler.poll_uex_notifications()
    outbox = memory.inspect()["outbox"]

    assert len(first) == 1
    assert second == []
    assert len(outbox) == 1
    assert "A buyer replied to your listing." in outbox[0]["content"]


@pytest.mark.asyncio
async def test_poll_uex_notifications_reports_failures_to_outbox(tmp_path):
    memory = MemoryStore(str(tmp_path / "memory.sqlite3"))
    scheduler = WakeScheduler(memory)
    scheduler.bind_uex_notifications(FailingUEXNotifications())

    result = await scheduler.poll_uex_notifications()

    assert result == []
    assert "bad token" in memory.inspect()["outbox"][0]["content"]


@pytest.mark.asyncio
async def test_wake_job_writes_agent_output_to_outbox_and_disables_one_shot(tmp_path):
    memory = MemoryStore(str(tmp_path / "memory.sqlite3"))
    scheduler = WakeScheduler(memory)
    scheduler.bind_agent(FakeWakeAgent())
    memory.add_job("wake-test", "check notifications", "date", "2099-01-01T00:00:00+00:00")

    await scheduler._run_job("wake-test", "check notifications")
    snapshot = memory.inspect()

    assert "Wake output:" in snapshot["outbox"][0]["content"]
    assert snapshot["scheduled_jobs"][0]["enabled"] == 0


@pytest.mark.asyncio
async def test_wake_job_checks_listings_and_writes_analysis_to_outbox(tmp_path):
    memory = MemoryStore(str(tmp_path / "memory.sqlite3"))
    scheduler = WakeScheduler(memory)
    agent = ListingWakeAgent(memory)
    scheduler.bind_agent(agent)
    memory.add_job("wake-listings", "check Wikelo Favor listings and analyze the market", "date", "2099-01-01T00:00:00+00:00")

    await scheduler._run_job("wake-listings", "check Wikelo Favor listings and analyze the market")
    snapshot = memory.inspect()
    content = snapshot["outbox"][0]["content"]

    assert agent.listing_tools.calls == [
        (
            "search_marketplace_listings",
            {"query": "Wikelo Favor", "operation": "sell", "limit": 5},
        )
    ]
    assert "Listing check complete" in content
    assert "500,000,000 UEC" in content
    assert "Suggested next action" in content