tests/test_pipeline.py

import unittest
from unittest.mock import patch

import pandas as pd
import torch

from utils import pipeline


class FakeEmbeddingModel:
    def __init__(self, vectors):
        self.vectors = vectors
        self.calls = []

    def __call__(self, text, pooling='mean'):
        self.calls.append((text, pooling))
        return torch.tensor(self.vectors[text])


class PipelineFactoryTest(unittest.TestCase):
    def test_originality_uses_embedding_factory(self):
        model = FakeEmbeddingModel(
            {
                'prompt': [1.0, 0.0],
                'response': [0.0, 1.0],
            }
        )
        df = pd.DataFrame({'prompt': ['prompt'], 'response': ['response']})

        with patch.object(pipeline, 'get_embedding_model', return_value=model) as factory:
            result = pipeline.p0_originality(df, 'fake-model', 'mean')

        factory.assert_called_once_with('fake-model')
        self.assertAlmostEqual(result.loc[0, 'originality'], 1.0)
        self.assertEqual(model.calls, [('prompt', 'mean'), ('response', 'mean')])

    def test_flexibility_uses_embedding_factory(self):
        model = FakeEmbeddingModel(
            {
                'p': [1.0, 0.0],
                'a': [1.0, 0.0],
                'b': [0.0, 1.0],
            }
        )
        df = pd.DataFrame(
            {
                'id': [1, 1, 1],
                'prompt': ['p', 'p', 'p'],
                'response': ['a', 'b', 'a'],
            }
        )

        with patch.object(pipeline, 'get_embedding_model', return_value=model) as factory:
            result = pipeline.p1_flexibility(df, 'fake-model', 'cls')

        factory.assert_called_once_with('fake-model')
        self.assertEqual(len(result), 1)
        self.assertAlmostEqual(result.loc[0, 'flexibility'], 2.0)
        self.assertEqual(model.calls, [('a', 'cls'), ('b', 'cls'), ('a', 'cls')])


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