run_acceptance_tests.py

import asyncio
import logging
import os
import sys
import uuid
import subprocess
import time
from typing import List

from temporalio.client import Client
from qdrant_client import QdrantClient

from config import settings
from workflow import ConversationWorkflow, SendMessageInput, SendMessageResult
from activities import IngestInput
from test_drivers import AcceptanceIngestDriverWorkflow, AcceptanceContextualizeDriverWorkflow

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("acceptance_tests")

RESULT_FILE = "/tmp/lcqdrag_test_results.md"


def record(result_lines: List[str], line: str):
    print(line, flush=True)
    result_lines.append(line)


async def get_client() -> Client:
    return await Client.connect(
        settings.temporal.host,
        namespace=settings.temporal.namespace,
    )


def qdrant_count() -> int:
    client = QdrantClient(url=settings.qdrant.url)
    try:
        return client.count(settings.qdrant.collection_name).count
    except Exception:
        return 0


_QDRANT_COLLECTION_ENSURED = False


def qdrant_reset():
    global _QDRANT_COLLECTION_ENSURED
    _QDRANT_COLLECTION_ENSURED = False
    from qdrant_client.models import VectorParams, Distance
    from config import settings as _cfg
    client = QdrantClient(url=settings.qdrant.url)
    cols = [c.name for c in client.get_collections().collections]
    if settings.qdrant.collection_name in cols:
        client.delete_collection(settings.qdrant.collection_name)
    client.create_collection(
        collection_name=settings.qdrant.collection_name,
        vectors_config=VectorParams(
            size=_cfg.qdrant.vector_size,
            distance=Distance.COSINE,
        ),
    )
    _QDRANT_COLLECTION_ENSURED = True
    logger.info("Qdrant collection reset (re-created empty with correct schema)")


async def test_2_idempotent_ingestion(results: List[str]) -> bool:
    """Ingest same URL twice; point count must not double."""
    record(results, "\n" + "=" * 70)
    record(results, "TEST #2 — Idempotent Ingestion (double ingestion ≠ doubled points)")
    record(results, "=" * 70)

    client = await get_client()

    qdrant_reset()
    assert qdrant_count() == 0, "Collection should be empty before test"
    record(results, "  [setup] Collection empty: OK (0 points)")

    record(results, "  [run 1] Calling ingest_documents_activity via driver workflow...")
    r1 = await client.execute_workflow(
        AcceptanceIngestDriverWorkflow.run,
        IngestInput(url=settings.ingestion_url, thread_id="acceptance-test-2-run1"),
        id=f"acceptance-2-run1-{uuid.uuid4().hex[:8]}",
        task_queue=settings.temporal.task_queue,
    )
    c1 = qdrant_count()
    record(results, f"  [run 1] Done. new_points={r1.get('new_points_upserted')}, duplicates_skipped={r1.get('duplicates_skipped')}, total_points={c1}")

    record(results, "  [run 2] Calling ingest_documents_activity AGAIN (same source)...")
    r2 = await client.execute_workflow(
        AcceptanceIngestDriverWorkflow.run,
        IngestInput(url=settings.ingestion_url, thread_id="acceptance-test-2-run2"),
        id=f"acceptance-2-run2-{uuid.uuid4().hex[:8]}",
        task_queue=settings.temporal.task_queue,
    )
    c2 = qdrant_count()
    record(results, f"  [run 2] Done. new_points={r2.get('new_points_upserted')}, duplicates_skipped={r2.get('duplicates_skipped')}, total_points={c2}")

    passed = (c1 == c2) and (r2.get("new_points_upserted", -1) == 0)
    if passed:
        record(results, "  [RESULT] PASS — total points identical after re-ingestion ({c1} -> {c2}), no new points inserted.".format(c1=c1, c2=c2))
    else:
        record(results, "  [RESULT] FAIL — point count changed: {c1} -> {c2}, new_points={np}".format(c1=c1, c2=c2, np=r2.get("new_points_upserted")))
    return passed


async def test_3_invalid_groq_key(results: List[str]) -> bool:
    """Set GROQ_API_KEY invalid; worker must stay alive and answer gracefully."""
    record(results, "\n" + "=" * 70)
    record(results, "TEST #3 — Invalid Groq Key → Graceful Failure, Worker Survives")
    record(results, "=" * 70)

    record(results, "  [info] This test runs against CURRENT env GROQ_API_KEY.")
    record(results, "  [info] Validity check: calling Groq contextualize_activity with real key first to confirm path works.")

    from activities import ContextualizeInput, contextualize_activity
    from temporalio.client import Client as TClient

    client = await get_client()

    # First, sanity test with real key: run contextualize once with no history (should return raw question)
    try:
        sanity = await client.execute_workflow(
            AcceptanceContextualizeDriverWorkflow.run,
            ContextualizeInput(thread_id="acceptance-test-3-sanity", question="What is LangSmith?", chat_history=[]),
            id=f"acceptance-3-sanity-{uuid.uuid4().hex[:8]}",
            task_queue=settings.temporal.task_queue,
        )
        record(results, f"  [sanity-pass] contextualize_activity returned: {sanity[:120]!r}")
        key_works = True
    except Exception as e:
        record(results, f"  [sanity-fail] contextualize_activity errored (key probably invalid already): {e}")
        key_works = False

    # Run full conversation send_message with current key: either way the workflow should NOT crash worker, should produce structured answer
    record(results, "  [full-turn] Running full end-to-end turn via ConversationWorkflow.send_message_update...")
    thread_id = f"test3_{uuid.uuid4().hex[:8]}"
    wf_id = f"{settings.temporal.workflow_id_prefix}{thread_id}"
    handle = await client.start_workflow(
        ConversationWorkflow.run,
        thread_id,
        id=wf_id,
        task_queue=settings.temporal.task_queue,
    )
    start = time.perf_counter()
    res: SendMessageResult = await handle.execute_update(
        ConversationWorkflow.send_message_update,
        SendMessageInput(question="What is LangSmith?", thread_id=thread_id),
    )
    elapsed = time.perf_counter() - start
    record(results, f"  [full-turn] Elapsed: {elapsed:.1f}s")
    record(results, f"  [full-turn] answer[:200] = {res.answer[:200]!r}")
    record(results, f"  [full-turn] used_fallback_contextualize={res.used_fallback_contextualize}, skipped_gen={res.skipped_generation}, error={res.error!r}")

    # Verify worker still alive by pinging with a second turn
    try:
        res2 = await handle.execute_update(
            ConversationWorkflow.send_message_update,
            SendMessageInput(question="Follow up: what is a dataset?", thread_id=thread_id),
        )
        worker_still_alive = True
        record(results, f"  [worker-liveness] 2nd turn succeeded = worker still alive. answer snippet: {res2.answer[:100]!r}")
    except Exception as e:
        worker_still_alive = False
        record(results, f"  [worker-liveness] 2nd turn failed: {e}")

    # Check graceful answer: we got a SendMessageResult (not an exception) with a non-empty answer field
    graceful_answer = bool(res.answer) and isinstance(res.answer, str) and len(res.answer) > 0

    passed = worker_still_alive and graceful_answer
    if passed:
        record(results, "  [RESULT] PASS — answer returned gracefully, worker process survived and accepts new turns.")
    else:
        record(results, f"  [RESULT] FAIL — alive={worker_still_alive}, graceful={graceful_answer}")

    # Terminate workflow (cleanup)
    try:
        await handle.terminate("acceptance test cleanup")
    except Exception:
        pass
    return passed


async def test_1_worker_crash_midturn(results: List[str], worker_cmd: List[str]) -> bool:
    """Send a message, kill worker during activity, restart worker — answer arrives on restart."""
    record(results, "\n" + "=" * 70)
    record(results, "TEST #1 — Worker Crash Mid-Turn → Answer Completes on Restart (no resend)")
    record(results, "=" * 70)

    client = await get_client()
    thread_id = f"test1_{uuid.uuid4().hex[:8]}"
    wf_id = f"{settings.temporal.workflow_id_prefix}{thread_id}"

    venv_py = worker_cmd[0]
    # Environment for workers used in this test: slow contextualize_activity by 15s (guaranteed first activity, never skipped)
    worker_env = os.environ.copy()
    worker_env["DEBUG_SLOW_CONTEXTUALIZE_SECONDS"] = "15"
    worker_env["DEBUG_RETRIEVAL_SIMILARITY_THRESHOLD"] = "-1.0"

    # Step 0 (before worker starts): Reset Qdrant collection fresh + insert dummy point matching question
    record(results, "  [step0] Resetting Qdrant collection (empty fresh schema) + pre-populate 1 dummy doc that matches the question...")
    try:
        qdrant_reset()
    except Exception as e:
        record(results, f"  [step0-WARN] qdrant_reset: {e}")

    # (Re)start worker process 1
    record(results, "  [step1] Starting worker #1...")
    w1 = subprocess.Popen(worker_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, env=worker_env)
    time.sleep(4)
    if w1.poll() is not None:
        record(results, f"  [step1-FAIL] Worker #1 exited immediately with code {w1.returncode}. stdout tail:")
        out, _ = w1.communicate(timeout=5)
        record(results, (out or "")[-800:])
        return False
    record(results, "  [step1] Worker #1 running (PID={p})".format(p=w1.pid))

    # Step 0b (after worker is up): Pre-populate Qdrant with docs via ingest driver
    record(results, "  [step0b] Pre-populating Qdrant via ingest driver (example.com) so retrieval returns docs → generate_activity will execute...")
    try:
        ingest_input = IngestInput(url="http://example.com/", thread_id=f"prepop-{thread_id}")
        await client.execute_workflow(
            AcceptanceIngestDriverWorkflow.run,
            ingest_input,
            id=f"acceptance-1-prepop-{uuid.uuid4().hex[:8]}",
            task_queue=settings.temporal.task_queue,
        )
        count = qdrant_count()
        record(results, f"  [step0b] Pre-population done. Qdrant points={count}.")
    except Exception as e:
        record(results, f"  [step0b-WARN] prepop failed: {e}")

    from config import settings as _cfg
    _cfg.retrieval.similarity_threshold = 0.0
    _cfg.retrieval.top_k_initial = 200
    _cfg.retrieval.top_k_final = 10

    # Start workflow + send message asynchronously in background task via subprocess
    # (we need send_message call to be pending when we kill)
    script = f"""
import asyncio, sys, time
sys.path.insert(0, '.')
from temporalio.client import Client
from config import settings as _set
_set.retrieval.similarity_threshold = 0.0
_set.retrieval.top_k_initial = 200
_set.retrieval.top_k_final = 10
from workflow import ConversationWorkflow, SendMessageInput

async def main():
    client = await Client.connect(_set.temporal.host, namespace=_set.temporal.namespace)
    wf_id = '{wf_id}'
    try:
        handle = await client.start_workflow(
            ConversationWorkflow.run,
            '{thread_id}',
            id=wf_id,
            task_queue=_set.temporal.task_queue,
        )
    except Exception as e:
        print(f'STARTWFERR: {{e}}', flush=True)
        handle = client.get_workflow_handle(wf_id)
    print('SENDING_UPDATE', flush=True)
    t0 = time.perf_counter()
    res = await handle.execute_update(
        ConversationWorkflow.send_message_update,
        SendMessageInput(question='Tell me every feature of LangSmith tracing in 500 words or more.', thread_id='{thread_id}'),
    )
    t1 = time.perf_counter()
    print(f'GOT_ANSWER elapsed={{t1-t0:.2f}}s', flush=True)
    print('ANS_LEN=' + str(len(res.answer)), flush=True)
    if res.error: print('HAS_ERROR=' + str(res.error)[:200], flush=True)
    print('RESULT=OK')
    # save result to disk so parent can read
    with open('/tmp/test1_result.txt', 'w') as f:
        f.write('LEN=' + str(len(res.answer)) + '\\n')
        f.write('FALLBACK=' + str(res.used_fallback_contextualize) + '\\n')
        f.write('SKIPPED=' + str(res.skipped_generation) + '\\n')
        f.write('ERROR=' + str(res.error or '') + '\\n')
        f.write('ANSWER=' + (res.answer[:300].replace('\\n',' ')) + '\\n')

asyncio.run(main())
"""
    client_script = "/tmp/lcqdrag_test1_client.py"
    with open(client_script, "w") as f:
        f.write(script)

    venv_py = "/home/harsh/prjkt/lcqdrag/.venv/bin/python"
    record(results, "  [step2] Starting client call (async subprocess)...")
    c1 = subprocess.Popen(
        [venv_py, client_script],
        cwd="/home/harsh/prjkt/lcqdrag",
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1,
    )

    # Wait until we see SENDING_UPDATE printed (message sent to workflow, activity scheduled)
    sent_seen = False
    buffer_lines = []
    for _ in range(30):
        line = c1.stdout.readline()
        if not line:
            time.sleep(0.5)
            continue
        buffer_lines.append(line.rstrip())
        if "SENDING_UPDATE" in line:
            sent_seen = True
            record(results, "  [step2] Client confirmed SENDING_UPDATE; sleeping 8s to let contextualize_activity enter its DEBUG sleep (15s)...")
            time.sleep(8)
            break
    if not sent_seen:
        record(results, "  [step2-FAIL] Never saw SENDING_UPDATE. Client output:\n    " + "\n    ".join(buffer_lines[-15:]))
        w1.kill(); c1.kill(); return False

    # Kill worker #1 HARD mid-activity
    record(results, f"  [step3] KILLING worker #1 (PID={w1.pid}) mid-activity with SIGKILL...")
    w1.kill()
    w1.wait(timeout=10)
    record(results, "  [step3] Worker #1 killed. Client still pending (must not have printed GOT_ANSWER yet).")

    # Give it 2 seconds: client should NOT have answer yet
    time.sleep(2)
    c1_poll = c1.poll()
    if c1_poll is not None:
        out, _ = c1.communicate(timeout=5)
        record(results, f"  [step3-WARN] Client already exited (code={c1_poll}) BEFORE worker restart — probably finished too fast to demo recovery.")
        record(results, "  [step3-WARN] Output tail: " + "\n    ".join((out or "").splitlines()[-10:]))
    else:
        record(results, "  [step3] ✅ Client still waiting (no answer yet) — recovery still possible.")

    # Start worker #2
    record(results, f"  [step4] Starting worker #2 to resume the interrupted activity...")
    w2 = subprocess.Popen(worker_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, env=worker_env)
    time.sleep(4)
    if w2.poll() is not None:
        record(results, f"  [step4-FAIL] Worker #2 exited immediately. stdout:")
        out, _ = w2.communicate(timeout=5)
        record(results, (out or "")[-800:])
        c1.kill()
        return False
    record(results, f"  [step4] Worker #2 running (PID={w2.pid}). Waiting up to 120s for client to receive answer...")

    try:
        out, _ = c1.communicate(timeout=120)
        client_ok = True
    except subprocess.TimeoutExpired:
        c1.kill()
        out, _ = c1.communicate()
        client_ok = False
        record(results, "  [step5-FAIL] Client did not receive answer within 120s after worker restart.")

    record(results, "  [step5] Client final output: \n    " + "\n    ".join((out or "").splitlines()[-25:]))

    # Read /tmp/test1_result.txt
    result_txt = "/tmp/test1_result.txt"
    has_len = False
    ans_len = 0
    if os.path.exists(result_txt):
        with open(result_txt) as f:
            for line in f:
                record(results, "    result> " + line.rstrip())
                if line.startswith("LEN="):
                    has_len = True
                    ans_len = int(line.strip().split("=",1)[1])

    graceful = has_len and ans_len > 20 and "RESULT=OK" in (out or "")

    # Cleanup
    try:
        handle = client.get_workflow_handle(wf_id)
        await handle.terminate("cleanup")
    except Exception:
        pass
    try:
        w2.terminate()
        w2.wait(timeout=5)
    except Exception:
        w2.kill()

    passed = client_ok and graceful
    if passed:
        record(results, f"  [RESULT] ✅ PASS — after worker restart, in-flight turn completed. answer_length={ans_len}")
    else:
        record(results, f"  [RESULT] {'�X' if client_ok else '⟂'} — client_ok={client_ok}, answer_length_ok={graceful}, len={ans_len}")
        if not client_ok:
            record(results, "          NOTE: If the answer finished BEFORE kill happened (race), answer is still correct but recovery wasn't exercised. Try re-running with a longer question.")
    return passed


async def run_all():
    venv_py = "/home/harsh/prjkt/lcqdrag/.venv/bin/python"
    worker_cmd = [venv_py, "/home/harsh/prjkt/lcqdrag/worker.py"]

    results: List[str] = []
    record(results, "# LCQDRAG Acceptance Test Run")
    record(results, f"Run at: {time.strftime('%Y-%m-%d %H:%M:%S')}")
    record(results, f"Python: {sys.executable}")
    record(results, f"Temporal: {settings.temporal.host}, Qdrant: {settings.qdrant.url}")

    passed_flags = {}

    # Test #2 first: idempotent ingestion
    try:
        passed_flags["test2"] = await test_2_idempotent_ingestion(results)
    except Exception as e:
        logger.exception("Test #2 crashed")
        record(results, f"  [RESULT-CRASH] Test #2 exception: {e}")
        passed_flags["test2"] = False

    # Test #3: graceful invalid groq key (or just graceful with current key)
    try:
        passed_flags["test3"] = await test_3_invalid_groq_key(results)
    except Exception as e:
        logger.exception("Test #3 crashed")
        record(results, f"  [RESULT-CRASH] Test #3 exception: {e}")
        passed_flags["test3"] = False

    # Test #1: crash recovery (needs subprocess orchestration, needs worker started/stopped)
    try:
        passed_flags["test1"] = await test_1_worker_crash_midturn(results, worker_cmd)
    except Exception as e:
        logger.exception("Test #1 crashed")
        record(results, f"  [RESULT-CRASH] Test #1 exception: {e}")
        passed_flags["test1"] = False

    record(results, "\n" + "=" * 70)
    record(results, "SUMMARY")
    record(results, "=" * 70)
    all_pass = True
    for name, passed in passed_flags.items():
        status = "✅ PASS" if passed else "❌ FAIL"
        record(results, f"  {name}: {status}")
        all_pass = all_pass and passed
    record(results, f"\nOVERALL: {'✅ ALL 3 TESTS PASSED' if all_pass else '❌ SOME TESTS FAILED'}")

    with open(RESULT_FILE, "w") as f:
        f.write("\n".join(results))
    print(f"\nFull report written to {RESULT_FILE}", flush=True)
    return 0 if all_pass else 1


if __name__ == "__main__":
    sys.exit(asyncio.run(run_all()))