scripts/test_silicon_logs.py

#!/usr/bin/env python3
"""Unit tests for anonymous silicon sampling run logs."""

from __future__ import annotations

import os
import sqlite3
import sys
import tempfile
import unittest
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from hf_dataset_logs import hf_log_path  # noqa: E402
from silicon_logs import log_status, record_run, recent_runs  # noqa: E402


class SiliconLogTests(unittest.TestCase):
    def test_record_run_writes_anonymous_settings_questions_and_summary(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "logs.sqlite3"
            old = os.environ.get("SILICON_LOG_DB")
            old_hf_token = os.environ.get("HF_TOKEN")
            old_hf_log_token = os.environ.get("HF_LOG_TOKEN")
            os.environ["SILICON_LOG_DB"] = str(path)
            os.environ.pop("HF_TOKEN", None)
            os.environ.pop("HF_LOG_TOKEN", None)
            try:
                info = record_run(
                    {
                        "config": {
                            "sampleSize": 2,
                            "genders": [{"id": "male", "enabled": True, "weight": 1}],
                            "ages": [{"id": "30s", "enabled": True, "weight": 2}],
                            "locations": [{"id": "seoul", "enabled": True, "weight": 2}],
                            "locationOptions": [{"id": "seoul", "label": "서울", "parentRegion": "seoul", "level": "sido"}],
                            "personaAttributes": [{"id": "occ_office", "enabled": True, "weight": 2}],
                            "nemotronFields": ["persona", "occupation"],
                            "questions": [
                                {"id": "q1", "title": "만족도", "kind": "likert", "scale": 5},
                                {"id": "q2", "title": "중요 가치", "kind": "categorical", "options": [{"id": "health", "label": "건강"}]},
                            ],
                        }
                    },
                    {
                        "respondents": [{"id": "R0001"}, {"id": "R0002"}],
                        "likertAnswers": [{"questionId": "q1"}, {"questionId": "q1"}],
                        "categoricalAnswers": [{"questionId": "q2"}, {"questionId": "q2"}],
                        "openAnswers": [],
                        "questionStats": [{"questionId": "q1"}, {"questionId": "q2"}],
                        "regionStats": [],
                        "llmTrace": {"calls": 2, "rawResponses": ["not stored"]},
                    },
                    client={"userAgent": "unit-test"},
                )
                self.assertTrue(info["runId"].startswith("run_"))
                self.assertEqual(info["hfDataset"], {"enabled": False, "uploaded": False})
                self.assertTrue(path.exists())
                self.assertEqual(log_status()["runCount"], 1)
                self.assertIn("hfDataset", log_status())
                rows = recent_runs()
                self.assertEqual(len(rows), 1)
                self.assertEqual(rows[0]["sampleSize"], 2)
                self.assertEqual(rows[0]["questions"][1]["kind"], "categorical")
                self.assertEqual(rows[0]["resultSummary"]["categoricalAnswers"], 2)
                self.assertNotIn("rawResponses", rows[0]["resultSummary"]["llmTrace"])
                with sqlite3.connect(path) as conn:
                    self.assertEqual(conn.execute("SELECT COUNT(*) FROM runs").fetchone()[0], 1)
            finally:
                if old is None:
                    os.environ.pop("SILICON_LOG_DB", None)
                else:
                    os.environ["SILICON_LOG_DB"] = old
                if old_hf_token is not None:
                    os.environ["HF_TOKEN"] = old_hf_token
                if old_hf_log_token is not None:
                    os.environ["HF_LOG_TOKEN"] = old_hf_log_token

    def test_record_run_can_mirror_json_document_to_hf_dataset_uploader(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "logs.sqlite3"
            old = os.environ.get("SILICON_LOG_DB")
            os.environ["SILICON_LOG_DB"] = str(path)
            uploaded: list[dict] = []

            def fake_uploader(document: dict) -> dict:
                uploaded.append(document)
                return {"enabled": True, "uploaded": True, "repoId": "kyu823/silicon-sampling-logs", "path": hf_log_path(document)}

            try:
                info = record_run(
                    {
                        "config": {
                            "sampleSize": 1,
                            "genders": [{"id": "female", "enabled": True, "weight": 1}],
                            "ages": [{"id": "40s", "enabled": True, "weight": 1}],
                            "locations": [{"id": "busan", "enabled": True, "weight": 1}],
                            "locationOptions": [{"id": "busan", "label": "부산", "parentRegion": "busan", "level": "sido"}],
                            "questions": [{"id": "q_cat", "title": "선택 문항", "kind": "categorical", "options": [{"id": "a", "label": "A"}]}],
                        }
                    },
                    {
                        "respondents": [{"id": "R0001"}],
                        "likertAnswers": [],
                        "categoricalAnswers": [{"questionId": "q_cat"}],
                        "openAnswers": [],
                        "questionStats": [{"questionId": "q_cat"}],
                        "regionStats": [],
                        "llmTrace": {"calls": 1, "rawResponses": ["not mirrored"]},
                    },
                    hf_uploader=fake_uploader,
                )
                self.assertEqual(info["hfDataset"]["repoId"], "kyu823/silicon-sampling-logs")
                self.assertTrue(info["hfDataset"]["path"].startswith("runs/"))
                self.assertEqual(len(uploaded), 1)
                self.assertEqual(uploaded[0]["questions"][0]["kind"], "categorical")
                self.assertEqual(uploaded[0]["respondentSettings"]["locations"][0]["label"], "부산")
                self.assertNotIn("rawResponses", uploaded[0]["resultSummary"]["llmTrace"])
            finally:
                if old is None:
                    os.environ.pop("SILICON_LOG_DB", None)
                else:
                    os.environ["SILICON_LOG_DB"] = old


if __name__ == "__main__":
    unittest.main(verbosity=2)