multiq/tests/test_difficulty.py

from unittest.mock import MagicMock

import pytest

from ..difficulty import DifficultyScorer


@pytest.fixture
def settings(tmp_path):
    from ..config import Settings
    return Settings(index_path=tmp_path / "idx.faiss")


class TestDifficulty:
    def test_easy_for_simple_overlap(self, settings):
        encoder = MagicMock()
        encoder.encode = lambda x, **kw: [[0.9] * 16 for _ in x]
        scorer = DifficultyScorer(settings, encoder)
        result = scorer.score(
            "Capital of France?",
            "Paris",
            ["London", "Berlin"],
            "The capital of France is Paris. It is known as the City of Light.",
        )
        assert result in ("easy", "medium", "hard")

    def test_hard_when_distractors_close(self, settings):
        encoder = MagicMock()
        # make distractor very similar to answer
        vectors = [
            [0.1, 0.9],  # context
            [0.15, 0.85],  # answer
            [0.16, 0.84],  # distractor 1
            [0.2, 0.8],  # distractor 2
        ]
        def fake_encode(x, **kw):
            return vectors[:len(x)]
        encoder.encode = fake_encode
        scorer = DifficultyScorer(settings, encoder)
        result = scorer.score(
            "Q",
            "Answer A",
            ["Answer A1", "Answer A2"],
            "Context text here",
        )
        assert result == "hard"