tests/conftest.py

from __future__ import annotations

import uuid
from datetime import date

import pytest

from civicsetu.models.enums import ChunkStatus, DocType, Jurisdiction, QueryType
from civicsetu.models.schemas import Citation, ChatMessage, LegalChunk, RetrievedChunk


def _make_chunk(
    section_id: str = "18",
    doc_name: str = "RERA Act 2016",
    jurisdiction: Jurisdiction = Jurisdiction.CENTRAL,
    text: str = "Section text placeholder.",
    doc_id: uuid.UUID | None = None,
) -> LegalChunk:
    return LegalChunk(
        chunk_id=uuid.uuid4(),
        doc_id=doc_id or uuid.uuid4(),
        jurisdiction=jurisdiction,
        doc_type=DocType.ACT,
        doc_name=doc_name,
        section_id=section_id,
        section_title=f"Section {section_id} Title",
        section_hierarchy=[doc_name, section_id],
        text=text,
        effective_date=date(2016, 5, 1),
        status=ChunkStatus.ACTIVE,
        source_url="https://example.com/rera.pdf",
        page_number=1,
    )


def _make_rc(
    section_id: str = "18",
    doc_name: str = "RERA Act 2016",
    jurisdiction: Jurisdiction = Jurisdiction.CENTRAL,
    is_pinned: bool = False,
    retrieval_source: str = "vector",
    rerank_score: float | None = None,
) -> RetrievedChunk:
    return RetrievedChunk(
        chunk=_make_chunk(section_id=section_id, doc_name=doc_name, jurisdiction=jurisdiction),
        retrieval_source=retrieval_source,
        is_pinned=is_pinned,
        rerank_score=rerank_score,
    )


def _base_state(**overrides) -> dict:
    state = {
        "query": "What are promoter obligations?",
        "session_id": "test-session",
        "jurisdiction_filter": None,
        "top_k": 5,
        "messages": [ChatMessage(role="user", content="What are promoter obligations?")],
        "query_type": None,
        "rewritten_query": None,
        "retrieved_chunks": [],
        "reranked_chunks": [],
        "raw_response": None,
        "citations": [],
        "confidence_score": 0.0,
        "conflict_warnings": [],
        "amendment_notice": None,
        "retry_count": 0,
        "hallucination_flag": False,
        "error": None,
    }
    state.update(overrides)
    return state


@pytest.fixture
def make_chunk():
    return _make_chunk


@pytest.fixture
def make_rc():
    return _make_rc


@pytest.fixture
def base_state():
    return _base_state