import unittest

from src.banned import banned_rule


class TestBanned(unittest.TestCase):
    def test_default_engine_and_openchat(self):
        for model, rule in [('org/Model-GGUF', 'gguf'), ('org/model.gguf', 'gguf'),
                            ('org/GGML-model', 'gguf'), ('OpenChat/model', 'openchat')]:
            with self.subTest(model=model):
                self.assertEqual(banned_rule(model), rule)

    def test_unmerged_adapter_hints(self):
        for model in ('org/model-lora', 'org/LORA', 'org/adapter_model',
                      'org/model-adapter-config', 'org/Model-PEFT'):
            with self.subTest(model=model):
                self.assertEqual(banned_rule(model), 'peft-lora')

    def test_full_weights_and_unrelated_names(self):
        for model in ('org/lora-merged', 'org/merge-lora', 'org/PEFT-FULLFT',
                      'org/full-ft-adapter_model', 'org/adapter-config-MERGED',
                      'org/flora', 'org/lorax', 'org/model_lora',
                      'org/model-GRPO', 'abaryan/CyberXP_Agent_Llama_3.2_1B'):
            with self.subTest(model=model):
                self.assertIsNone(banned_rule(model))

    def test_full_weight_hint_does_not_override_other_bans(self):
        self.assertEqual(banned_rule('org/lora-merged-GGUF'), 'gguf')
        self.assertEqual(banned_rule('org/openchat-fullft'), 'openchat')

    def test_existing_rules(self):
        for model in ('gpt2', 'org/Phi-3-mini', 'nvidia/Nemotron-H-4B',
                      'allenai/OLMo-2-7B', 'org/LFM2', 'org/Spark-X2.5',
                      'org/Breeze', 'org/gemma2', 'org/mistral-v0.1', 'org/tinyllama'):
            with self.subTest(model=model):
                self.assertIsNotNone(banned_rule(model))