anya/tests/test_aliases.py

import os
import tempfile
import unittest

from services.normalization import get_aliases, normalize_headline, replace_aliases


ALIASES = {
    "federal reserve": "fed",
    "united states": "us",
    "united states of america": "us",
    "president of the united states": "potus",
    "united nations": "unitednations",
    "inflation rate": "inflation",
}


class TestGetAliases(unittest.TestCase):
    def test_loads_toml_with_key_normalization(self):
        path = None
        try:
            with tempfile.NamedTemporaryFile("wb", suffix=".toml", delete=False) as f:
                f.write(
                    b"[aliases]\n"
                    b"\"Federal Reserve\" = \"Fed\"\n"
                    b"\"people's republic of china\" = \"china\"\n"
                )
                path = f.name
            aliases = get_aliases(path)
            self.assertEqual(aliases["federal reserve"], "fed")
            # apostrophe is stripped so keys match punctuation-removed text
            self.assertEqual(aliases["peoples republic of china"], "china")
        finally:
            if path:
                os.unlink(path)

    def test_missing_file_raises(self):
        with self.assertRaises(FileNotFoundError):
            get_aliases("/no/such/aliases.toml")


class TestReplaceAliases(unittest.TestCase):
    def test_whole_word_only(self):
        # "united states" must not match inside "united statesman".
        self.assertEqual(
            replace_aliases("united statesman wins award", ALIASES),
            "united statesman wins award",
        )

    def test_longest_phrase_first(self):
        # "president of the united states" -> "potus", not "president of the us".
        self.assertEqual(
            replace_aliases("president of the united states speaks", ALIASES),
            "potus speaks",
        )

    def test_no_aliases_returns_text(self):
        self.assertEqual(replace_aliases("plain text here", None), "plain text here")
        self.assertEqual(replace_aliases("plain text here", {}), "plain text here")


class TestNormalizeHeadlineAliases(unittest.TestCase):
    def setUp(self):
        self.stopwords = {"the", "a", "an", "in", "on", "and", "of", "to", "for"}

    def test_full_name_and_acronym_cluster(self):
        # "Federal Reserve" and "the Fed" normalize to the same first token.
        full = normalize_headline("Federal Reserve raises rates", self.stopwords, ALIASES)
        short = normalize_headline("the Fed raises rates", self.stopwords, ALIASES)
        self.assertEqual(full, ["fed", "raises", "rates"])
        self.assertEqual(short, ["fed", "raises", "rates"])

    def test_united_nations_not_dropped(self):
        # Regression guard: the canonical must survive stopword removal.
        self.assertEqual(
            normalize_headline("United Nations meets", self.stopwords, ALIASES),
            ["unitednations", "meets"],
        )

    def test_inflation_rates_not_mangled(self):
        # "inflation rate" must not fire inside "inflation rates".
        self.assertEqual(
            normalize_headline("inflation rates fall", self.stopwords, ALIASES),
            ["inflation", "rates", "fall"],
        )


class TestAliasesFileIntegrity(unittest.TestCase):
    """Guard against an alias canonical colliding with a stopword, which would
    get silently stripped during normalization (the original 'united nations'
    -> 'un' bug)."""

    def test_no_canonical_is_a_stopword(self):
        base = os.path.join(os.path.dirname(__file__), "..", "resources")
        aliases = get_aliases(os.path.join(base, "aliases.toml"))
        with open(os.path.join(base, "stopwords.txt"), encoding="utf-8-sig") as f:
            stopwords = set(f.read().splitlines())
        self.assertTrue(aliases)
        for phrase, canonical in aliases.items():
            for token in canonical.split():
                self.assertNotIn(
                    token,
                    stopwords,
                    "alias %r canonical token %r is a stopword" % (phrase, token),
                )


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