Mcp Server Dev

SQLite format 3@  
.�
�"��y"U#
indexidx_claimedtasksCREATE INDEX idx_claimed ON tasks(claimed_by, claimed_at)U)
yindexidx_job_statustasksCREATE INDEX idx_job_status ON tasks(job_name, status)�z

�StabletaskstasksCREATE TABLE tasks (
  id TEXT PRIMARY KEY,
  job_name TEXT NOT NULL,
  status TEXT NOT NULL DEFAULT 'pending',
  claimed_by TEXT,
  claimed_at TEXT,
  updated_at TEXT DEFAULT CURRENT_TIMESTAMP,
  reason TEXT,
  metadata TEXT
))=
indexsqlite_autoindex_tasks_1tasks





"""postgres-control-plane: pluggable task queue with FastMCP server."""
from .backend import (
    MemoryBackend,
    PostgresBackend,
    SQLiteBackend,
    TaskBackend,
    make_backend,
)
from .models import Task, TaskStatus, is_valid_transition
from .queue import TaskQueue

__all__ = [
    "Task",
    "TaskStatus",
    "TaskQueue",
    "TaskBackend",
    "MemoryBackend",
    "SQLiteBackend",
    "PostgresBackend",
    "make_backend",
    "is_valid_transition",
]

__version__ = "0.2.0"

__pycache__/
*.pyc
*.pyo
.tasks/
*.egg-info/
build/
dist/
.pytest_cache/

"""Pluggable task storage backends: Memory, SQLite, Postgres."""
from __future__ import annotations

import json
import os
import sqlite3
import threading
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Optional

from .models import Task, is_valid_transition


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


class TaskBackend(ABC):
    @abstractmethod
    def add(self, task: Task) -> Task: ...

    @abstractmethod
    def get(self, task_id: str) -> Optional[Task]: ...

    @abstractmethod
    def list_pending(self, job_name: str) -> List[Task]: ...

    @abstractmethod
    def claim(self, task_id: str, worker_id: str) -> Task: ...

    @abstractmethod
    def update_status(self, task_id: str, status: str, reason: str = "") -> Task: ...

    @abstractmethod
    def heartbeat(self, task_id: str, worker_id: str) -> Task: ...

    @abstractmethod
    def get_stats(self, job_name: str) -> dict: ...


# ---------------------------------------------------------------------------
# Memory
# ---------------------------------------------------------------------------
class MemoryBackend(TaskBackend):
    def __init__(self):
        self._tasks: dict[str, Task] = {}
        self._lock = threading.Lock()

    def add(self, task: Task) -> Task:
        with self._lock:
            task.updated_at = _now()
            self._tasks[task.id] = task
            return task

    def get(self, task_id: str) -> Optional[Task]:
        return self._tasks.get(task_id)

    def list_pending(self, job_name: str) -> List[Task]:
        return [t for t in self._tasks.values()
                if t.job_name == job_name and t.status == "pending"]

    def claim(self, task_id: str, worker_id: str) -> Task:
        with self._lock:
            task = self._tasks.get(task_id)
            if not task:
                raise KeyError(f"Task {task_id} not found")
            if task.status != "pending":
                raise ValueError(f"Task {task_id} not pending (status={task.status})")
            task.status = "claimed"
            task.claimed_by = worker_id
            task.claimed_at = _now()
            task.updated_at = task.claimed_at
            return task

    def update_status(self, task_id: str, status: str, reason: str = "") -> Task:
        with self._lock:
            task = self._tasks.get(task_id)
            if not task:
                raise KeyError(f"Task {task_id} not found")
            if not is_valid_transition(task.status, status):
                raise ValueError(f"Invalid transition: {task.status} -> {status}")
            task.status = status
            task.reason = reason
            task.updated_at = _now()
            return task

    def heartbeat(self, task_id: str, worker_id: str) -> Task:
        with self._lock:
            task = self._tasks.get(task_id)
            if not task:
                raise KeyError(f"Task {task_id} not found")
            if task.claimed_by != worker_id:
                raise ValueError(f"Task {task_id} not claimed by {worker_id}")
            task.claimed_at = _now()
            task.updated_at = task.claimed_at
            return task

    def get_stats(self, job_name: str) -> dict:
        tasks = [t for t in self._tasks.values() if t.job_name == job_name]
        return {
            "pending": sum(1 for t in tasks if t.status == "pending"),
            "claimed": sum(1 for t in tasks if t.status == "claimed"),
            "completed": sum(1 for t in tasks if t.status == "completed"),
            "failed": sum(1 for t in tasks if t.status == "failed"),
        }


# ---------------------------------------------------------------------------
# SQLite
# ---------------------------------------------------------------------------
SQLITE_SCHEMA = """
CREATE TABLE IF NOT EXISTS tasks (
  id TEXT PRIMARY KEY,
  job_name TEXT NOT NULL,
  status TEXT NOT NULL DEFAULT 'pending',
  claimed_by TEXT,
  claimed_at TEXT,
  updated_at TEXT DEFAULT CURRENT_TIMESTAMP,
  reason TEXT,
  metadata TEXT
);
CREATE INDEX IF NOT EXISTS idx_job_status ON tasks(job_name, status);
CREATE INDEX IF NOT EXISTS idx_claimed ON tasks(claimed_by, claimed_at);
"""


def _row_to_task(row) -> Task:
    return Task(
        id=row["id"],
        job_name=row["job_name"],
        status=row["status"],
        claimed_by=row["claimed_by"],
        claimed_at=row["claimed_at"],
        updated_at=row["updated_at"],
        reason=row["reason"],
        metadata=json.loads(row["metadata"]) if row["metadata"] else {},
    )


class SQLiteBackend(TaskBackend):
    def __init__(self, db_path: str = ".tasks/queue.db"):
        self.db_path = db_path
        Path(db_path).parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        with self._conn() as conn:
            conn.executescript(SQLITE_SCHEMA)

    def _conn(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path, isolation_level=None, timeout=30)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA foreign_keys=ON")
        return conn

    def add(self, task: Task) -> Task:
        task.updated_at = _now()
        with self._lock, self._conn() as conn:
            conn.execute(
                """INSERT INTO tasks(id, job_name, status, claimed_by, claimed_at,
                                     updated_at, reason, metadata)
                   VALUES(?, ?, ?, ?, ?, ?, ?, ?)""",
                (task.id, task.job_name, task.status, task.claimed_by,
                 task.claimed_at, task.updated_at, task.reason,
                 json.dumps(task.metadata or {})),
            )
        return task

    def get(self, task_id: str) -> Optional[Task]:
        with self._conn() as conn:
            row = conn.execute("SELECT * FROM tasks WHERE id=?", (task_id,)).fetchone()
        return _row_to_task(row) if row else None

    def list_pending(self, job_name: str) -> List[Task]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT * FROM tasks WHERE job_name=? AND status='pending'",
                (job_name,),
            ).fetchall()
        return [_row_to_task(r) for r in rows]

    def claim(self, task_id: str, worker_id: str) -> Task:
        ts = _now()
        with self._lock, self._conn() as conn:
            cur = conn.execute(
                """UPDATE tasks
                   SET status='claimed', claimed_by=?, claimed_at=?, updated_at=?
                   WHERE id=? AND status='pending'""",
                (worker_id, ts, ts, task_id),
            )
            if cur.rowcount == 0:
                row = conn.execute("SELECT status FROM tasks WHERE id=?", (task_id,)).fetchone()
                if not row:
                    raise KeyError(f"Task {task_id} not found")
                raise ValueError(f"Task {task_id} not pending (status={row['status']})")
        task = self.get(task_id)
        assert task is not None
        return task

    def update_status(self, task_id: str, status: str, reason: str = "") -> Task:
        with self._lock, self._conn() as conn:
            row = conn.execute("SELECT status FROM tasks WHERE id=?", (task_id,)).fetchone()
            if not row:
                raise KeyError(f"Task {task_id} not found")
            if not is_valid_transition(row["status"], status):
                raise ValueError(f"Invalid transition: {row['status']} -> {status}")
            conn.execute(
                "UPDATE tasks SET status=?, reason=?, updated_at=? WHERE id=?",
                (status, reason, _now(), task_id),
            )
        task = self.get(task_id)
        assert task is not None
        return task

    def heartbeat(self, task_id: str, worker_id: str) -> Task:
        ts = _now()
        with self._lock, self._conn() as conn:
            cur = conn.execute(
                "UPDATE tasks SET claimed_at=?, updated_at=? WHERE id=? AND claimed_by=?",
                (ts, ts, task_id, worker_id),
            )
            if cur.rowcount == 0:
                row = conn.execute("SELECT id FROM tasks WHERE id=?", (task_id,)).fetchone()
                if not row:
                    raise KeyError(f"Task {task_id} not found")
                raise ValueError(f"Task {task_id} not claimed by {worker_id}")
        task = self.get(task_id)
        assert task is not None
        return task

    def get_stats(self, job_name: str) -> dict:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT status, COUNT(*) AS n FROM tasks WHERE job_name=? GROUP BY status",
                (job_name,),
            ).fetchall()
        stats = {"pending": 0, "claimed": 0, "completed": 0, "failed": 0}
        for r in rows:
            stats[r["status"]] = r["n"]
        return stats


# ---------------------------------------------------------------------------
# Postgres
# ---------------------------------------------------------------------------
POSTGRES_SCHEMA = """
CREATE TABLE IF NOT EXISTS tasks (
  id TEXT PRIMARY KEY,
  job_name TEXT NOT NULL,
  status TEXT NOT NULL DEFAULT 'pending',
  claimed_by TEXT,
  claimed_at TIMESTAMPTZ,
  updated_at TIMESTAMPTZ DEFAULT NOW(),
  reason TEXT,
  metadata JSONB
);
CREATE INDEX IF NOT EXISTS idx_job_status ON tasks(job_name, status);
CREATE INDEX IF NOT EXISTS idx_claimed ON tasks(claimed_by, claimed_at);
"""


class PostgresBackend(TaskBackend):
    def __init__(self, dsn: str):
        try:
            import psycopg2
            from psycopg2.extras import RealDictCursor, Json
        except ImportError as e:
            raise RuntimeError(
                "psycopg2-binary required for PostgresBackend. "
                "Install: pip install psycopg2-binary"
            ) from e
        self._psycopg2 = psycopg2
        self._RealDictCursor = RealDictCursor
        self._Json = Json
        self.dsn = dsn
        with self._conn() as conn, conn.cursor() as cur:
            cur.execute(POSTGRES_SCHEMA)
            conn.commit()

    def _conn(self):
        return self._psycopg2.connect(self.dsn)

    def _row_to_task(self, row) -> Task:
        return Task(
            id=row["id"],
            job_name=row["job_name"],
            status=row["status"],
            claimed_by=row["claimed_by"],
            claimed_at=row["claimed_at"].isoformat() if row["claimed_at"] else None,
            updated_at=row["updated_at"].isoformat() if row["updated_at"] else None,
            reason=row["reason"],
            metadata=row["metadata"] or {},
        )

    def add(self, task: Task) -> Task:
        with self._conn() as conn, conn.cursor(cursor_factory=self._RealDictCursor) as cur:
            cur.execute(
                """INSERT INTO tasks(id, job_name, status, claimed_by, claimed_at,
                                     reason, metadata)
                   VALUES(%s, %s, %s, %s, %s, %s, %s)
                   RETURNING *""",
                (task.id, task.job_name, task.status, task.claimed_by,
                 task.claimed_at, task.reason, self._Json(task.metadata or {})),
            )
            row = cur.fetchone()
            conn.commit()
        return self._row_to_task(row)

    def get(self, task_id: str) -> Optional[Task]:
        with self._conn() as conn, conn.cursor(cursor_factory=self._RealDictCursor) as cur:
            cur.execute("SELECT * FROM tasks WHERE id=%s", (task_id,))
            row = cur.fetchone()
        return self._row_to_task(row) if row else None

    def list_pending(self, job_name: str) -> List[Task]:
        with self._conn() as conn, conn.cursor(cursor_factory=self._RealDictCursor) as cur:
            cur.execute(
                "SELECT * FROM tasks WHERE job_name=%s AND status='pending'",
                (job_name,),
            )
            rows = cur.fetchall()
        return [self._row_to_task(r) for r in rows]

    def claim(self, task_id: str, worker_id: str) -> Task:
        with self._conn() as conn, conn.cursor(cursor_factory=self._RealDictCursor) as cur:
            cur.execute(
                """UPDATE tasks
                   SET status='claimed', claimed_by=%s, claimed_at=NOW(), updated_at=NOW()
                   WHERE id=%s AND status='pending'
                   RETURNING *""",
                (worker_id, task_id),
            )
            row = cur.fetchone()
            if not row:
                cur.execute("SELECT status FROM tasks WHERE id=%s", (task_id,))
                check = cur.fetchone()
                conn.rollback()
                if not check:
                    raise KeyError(f"Task {task_id} not found")
                raise ValueError(f"Task {task_id} not pending (status={check['status']})")
            conn.commit()
        return self._row_to_task(row)

    def update_status(self, task_id: str, status: str, reason: str = "") -> Task:
        with self._conn() as conn, conn.cursor(cursor_factory=self._RealDictCursor) as cur:
            cur.execute("SELECT status FROM tasks WHERE id=%s", (task_id,))
            row = cur.fetchone()
            if not row:
                raise KeyError(f"Task {task_id} not found")
            if not is_valid_transition(row["status"], status):
                raise ValueError(f"Invalid transition: {row['status']} -> {status}")
            cur.execute(
                """UPDATE tasks SET status=%s, reason=%s, updated_at=NOW()
                   WHERE id=%s RETURNING *""",
                (status, reason, task_id),
            )
            updated = cur.fetchone()
            conn.commit()
        return self._row_to_task(updated)

    def heartbeat(self, task_id: str, worker_id: str) -> Task:
        with self._conn() as conn, conn.cursor(cursor_factory=self._RealDictCursor) as cur:
            cur.execute(
                """UPDATE tasks SET claimed_at=NOW(), updated_at=NOW()
                   WHERE id=%s AND claimed_by=%s
                   RETURNING *""",
                (task_id, worker_id),
            )
            row = cur.fetchone()
            if not row:
                cur.execute("SELECT id FROM tasks WHERE id=%s", (task_id,))
                check = cur.fetchone()
                conn.rollback()
                if not check:
                    raise KeyError(f"Task {task_id} not found")
                raise ValueError(f"Task {task_id} not claimed by {worker_id}")
            conn.commit()
        return self._row_to_task(row)

    def get_stats(self, job_name: str) -> dict:
        with self._conn() as conn, conn.cursor(cursor_factory=self._RealDictCursor) as cur:
            cur.execute(
                "SELECT status, COUNT(*) AS n FROM tasks WHERE job_name=%s GROUP BY status",
                (job_name,),
            )
            rows = cur.fetchall()
        stats = {"pending": 0, "claimed": 0, "completed": 0, "failed": 0}
        for r in rows:
            stats[r["status"]] = r["n"]
        return stats


# ---------------------------------------------------------------------------
# Auto-detect
# ---------------------------------------------------------------------------
def make_backend(backend: str = "auto", **kwargs) -> TaskBackend:
    """Construct a backend.

    backend: "memory" | "sqlite" | "postgres" | "auto"
    auto: DATABASE_URL env -> postgres, else .tasks/ dir present or creatable -> sqlite, else memory
    """
    if backend == "memory":
        return MemoryBackend()
    if backend == "sqlite":
        return SQLiteBackend(kwargs.get("db_path", ".tasks/queue.db"))
    if backend == "postgres":
        dsn = kwargs.get("dsn") or os.environ.get("DATABASE_URL")
        if not dsn:
            raise ValueError("PostgresBackend requires dsn or DATABASE_URL env")
        return PostgresBackend(dsn)
    if backend == "auto":
        dsn = os.environ.get("DATABASE_URL")
        if dsn:
            return PostgresBackend(dsn)
        try:
            return SQLiteBackend(".tasks/queue.db")
        except Exception:
            return MemoryBackend()
    raise ValueError(f"Unknown backend: {backend}")

"""pytest config: make postgres_control_plane importable when running tests from inside the package dir."""
import sys
from pathlib import Path

# Add mcp-server-dev/ (the parent of postgres_control_plane/) to sys.path
# so that `import postgres_control_plane` works.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

"""FastMCP server exposing TaskQueue as 7 tools.

Run:
    python -m postgres_control_plane.mcp_server

Backend selection via env:
    DATABASE_URL=postgresql://... python -m postgres_control_plane.mcp_server
"""
from __future__ import annotations

import os
from dataclasses import asdict

from fastmcp import FastMCP

from .models import Task, is_valid_transition
from .queue import TaskQueue

BACKEND = os.environ.get("CONTROL_PLANE_BACKEND", "auto")
_backend_kwargs: dict = {}
_db_path = os.environ.get("CONTROL_PLANE_DB_PATH")
if _db_path:
    _backend_kwargs["db_path"] = _db_path
_dsn = os.environ.get("DATABASE_URL")
if _dsn and BACKEND in ("postgres", "auto"):
    _backend_kwargs.setdefault("dsn", _dsn)
queue = TaskQueue(backend=BACKEND, **_backend_kwargs)

mcp = FastMCP("control-plane")


@mcp.tool()
def add_task(task_id: str, job_name: str, metadata: dict | None = None) -> dict:
    """Add a new task in pending status.

    Args:
        task_id: Unique identifier.
        job_name: Job grouping label (e.g., 'data-processing').
        metadata: Optional dict of task-specific data.
    """
    task = Task(id=task_id, job_name=job_name, metadata=metadata or {})
    saved = queue.add(task)
    return asdict(saved)


@mcp.tool()
def list_tasks(job_name: str, status: str = "pending") -> dict:
    """List tasks by job and status. Only 'pending' is currently indexed."""
    if status == "pending":
        tasks = queue.list_pending(job_name)
    else:
        # Fallback: filter from stats (no full enumeration API by design)
        tasks = []
    return {
        "job_name": job_name,
        "status": status,
        "count": len(tasks),
        "tasks": [asdict(t) for t in tasks],
    }


@mcp.tool()
def claim_task(task_id: str, worker_id: str) -> dict:
    """Atomically claim a pending task. Raises if already claimed."""
    task = queue.claim(task_id, worker_id)
    return asdict(task)


@mcp.tool()
def update_task_status(task_id: str, status: str, reason: str = "") -> dict:
    """Update task status. Validates transition (pending->claimed->completed/failed)."""
    task = queue.update_status(task_id, status, reason)
    return asdict(task)


@mcp.tool()
def heartbeat(task_id: str, worker_id: str) -> dict:
    """Refresh claim lease — updates claimed_at to NOW()."""
    task = queue.heartbeat(task_id, worker_id)
    return {"task_id": task.id, "heartbeat": "ok", "claimed_at": task.claimed_at}


@mcp.tool()
def verify_transition(from_status: str, to_status: str) -> dict:
    """Check if a status transition is valid (without mutating)."""
    valid = is_valid_transition(from_status, to_status)
    return {"from": from_status, "to": to_status, "valid": valid}


@mcp.tool()
def get_stats(job_name: str) -> dict:
    """Get task counts by status for a job."""
    stats = queue.get_stats(job_name)
    return {"job_name": job_name, "stats": stats, "total": sum(stats.values())}


def main():
    mcp.run()


if __name__ == "__main__":
    main()

"""Task model and status enum."""
from dataclasses import dataclass, field, asdict
from enum import Enum
from typing import Optional


class TaskStatus(str, Enum):
    PENDING = "pending"
    CLAIMED = "claimed"
    COMPLETED = "completed"
    FAILED = "failed"


VALID_TRANSITIONS = {
    "pending": {"claimed"},
    "claimed": {"completed", "failed", "pending"},
    "completed": set(),
    "failed": set(),
}


def is_valid_transition(from_status: str, to_status: str) -> bool:
    return to_status in VALID_TRANSITIONS.get(from_status, set())


@dataclass
class Task:
    id: str
    job_name: str
    status: str = TaskStatus.PENDING.value
    claimed_by: Optional[str] = None
    claimed_at: Optional[str] = None
    updated_at: Optional[str] = None
    reason: Optional[str] = None
    metadata: Optional[dict] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return asdict(self)

"""TaskQueue public API. Delegates to a pluggable backend."""
from __future__ import annotations

from typing import List, Optional

from .backend import TaskBackend, make_backend
from .models import Task


class TaskQueue:
    """Generic task queue with pluggable backend.

    Examples
    --------
    >>> q = TaskQueue()                       # auto-detect backend
    >>> q = TaskQueue(backend="memory")       # explicit
    >>> q = TaskQueue(backend="sqlite", db_path=".tasks/q.db")
    >>> q = TaskQueue(backend="postgres", dsn="postgresql://...")
    """

    def __init__(self, backend: str = "auto", **backend_kwargs):
        if isinstance(backend, TaskBackend):
            self.backend: TaskBackend = backend
        else:
            self.backend = make_backend(backend, **backend_kwargs)

    def add(self, task: Task) -> Task:
        return self.backend.add(task)

    def get(self, task_id: str) -> Optional[Task]:
        return self.backend.get(task_id)

    def list_pending(self, job_name: str) -> List[Task]:
        return self.backend.list_pending(job_name)

    def claim(self, task_id: str, worker_id: str) -> Task:
        return self.backend.claim(task_id, worker_id)

    def update_status(self, task_id: str, status: str, reason: str = "") -> Task:
        return self.backend.update_status(task_id, status, reason)

    def heartbeat(self, task_id: str, worker_id: str) -> Task:
        return self.backend.heartbeat(task_id, worker_id)

    def get_stats(self, job_name: str) -> dict:
        return self.backend.get_stats(job_name)

from setuptools import setup, find_packages

setup(
    name="postgres-control-plane",
    version="0.2.0",
    description="Pluggable task queue (Memory/SQLite/Postgres) with FastMCP server",
    packages=find_packages(),
    python_requires=">=3.10",
    install_requires=[
        "fastmcp>=0.1.0",
    ],
    extras_require={
        "postgres": ["psycopg2-binary>=2.9.0"],
        "dev": ["pytest>=7.0"],
    },
    entry_points={
        "console_scripts": [
            "control-plane=postgres_control_plane.mcp_server:main",
        ],
    },
)

"""Backend conformance + claim atomicity tests.

Run:
    pytest mcp-server-dev/postgres-control-plane/tests/

Postgres tests are skipped unless DATABASE_URL is set.
"""
from __future__ import annotations

import os
import threading
import tempfile

import pytest

from postgres_control_plane import (
    MemoryBackend,
    PostgresBackend,
    SQLiteBackend,
    Task,
    TaskQueue,
)


def make_queue_memory():
    return TaskQueue(backend=MemoryBackend())


def make_queue_sqlite():
    tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
    tmp.close()
    return TaskQueue(backend=SQLiteBackend(tmp.name))


BACKENDS = [make_queue_memory, make_queue_sqlite]

if os.environ.get("DATABASE_URL"):
    def make_queue_postgres():
        return TaskQueue(backend=PostgresBackend(os.environ["DATABASE_URL"]))
    BACKENDS.append(make_queue_postgres)


@pytest.fixture(params=BACKENDS, ids=lambda f: f.__name__)
def queue(request):
    return request.param()


def test_add_and_get(queue):
    queue.add(Task(id="t1", job_name="demo"))
    t = queue.get("t1")
    assert t and t.id == "t1" and t.status == "pending"


def test_list_pending(queue):
    queue.add(Task(id="a", job_name="demo"))
    queue.add(Task(id="b", job_name="demo"))
    queue.add(Task(id="c", job_name="other"))
    pending = queue.list_pending("demo")
    assert {t.id for t in pending} == {"a", "b"}


def test_claim_sets_metadata(queue):
    queue.add(Task(id="x", job_name="demo"))
    claimed = queue.claim("x", "worker-1")
    assert claimed.status == "claimed"
    assert claimed.claimed_by == "worker-1"
    assert claimed.claimed_at is not None
    assert claimed.updated_at is not None


def test_claim_twice_fails(queue):
    queue.add(Task(id="y", job_name="demo"))
    queue.claim("y", "w1")
    with pytest.raises(ValueError):
        queue.claim("y", "w2")


def test_claim_missing_fails(queue):
    with pytest.raises(KeyError):
        queue.claim("nope", "w1")


def test_status_transition_valid(queue):
    queue.add(Task(id="s1", job_name="demo"))
    queue.claim("s1", "w1")
    done = queue.update_status("s1", "completed", reason="ok")
    assert done.status == "completed"
    assert done.reason == "ok"


def test_status_transition_invalid(queue):
    queue.add(Task(id="s2", job_name="demo"))
    with pytest.raises(ValueError):
        queue.update_status("s2", "completed")  # pending -> completed not allowed


def test_heartbeat_updates_claimed_at(queue):
    queue.add(Task(id="h1", job_name="demo"))
    claimed = queue.claim("h1", "w1")
    first_ts = claimed.claimed_at
    import time; time.sleep(0.01)
    refreshed = queue.heartbeat("h1", "w1")
    assert refreshed.claimed_at != first_ts


def test_heartbeat_wrong_worker(queue):
    queue.add(Task(id="h2", job_name="demo"))
    queue.claim("h2", "w1")
    with pytest.raises(ValueError):
        queue.heartbeat("h2", "w2")


def test_stats(queue):
    queue.add(Task(id="g1", job_name="g"))
    queue.add(Task(id="g2", job_name="g"))
    queue.claim("g1", "w1")
    stats = queue.get_stats("g")
    assert stats["pending"] == 1
    assert stats["claimed"] == 1


def test_claim_atomicity(queue):
    """Race 10 workers against 1 task. Exactly one must win."""
    queue.add(Task(id="race", job_name="demo"))
    winners: list[str] = []
    lock = threading.Lock()

    def worker(wid: str):
        try:
            queue.claim("race", wid)
            with lock:
                winners.append(wid)
        except ValueError:
            pass

    threads = [threading.Thread(target=worker, args=(f"w{i}",)) for i in range(10)]
    for t in threads: t.start()
    for t in threads: t.join()

    assert len(winners) == 1, f"Expected 1 winner, got {len(winners)}: {winners}"