# tests/test_narrative_axis.py
from __future__ import annotations

import random
import unittest

from daily.narrative_axis import (
    NARRATIVE_AXES,
    enforce_narrative_axis,
    pick_narrative_axis,
)


class NarrativeAxisTests(unittest.TestCase):
    def test_pick_excludes_used(self):
        used = {"政策监管", "模型发布", "工具链/Agent"}
        for _ in range(20):
            axis = pick_narrative_axis(used, rng=random.Random(1))
            self.assertNotIn(axis, used)
            self.assertIn(axis, NARRATIVE_AXES)

    def test_enforce_overwrites_llm(self):
        trends = {"narrative_axis": "开源生态", "opening": "..."}
        out = enforce_narrative_axis(trends, "芯片算力")
        self.assertEqual(out["narrative_axis"], "芯片算力")

    def test_pick_when_all_used_falls_back(self):
        used = set(NARRATIVE_AXES)
        axis = pick_narrative_axis(used, rng=random.Random(0))
        self.assertIn(axis, NARRATIVE_AXES)


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