tests/test_app.py

import unittest

import app
from utils.models import QWEN3_EMBEDDING_MODEL


class PoolingUiTest(unittest.TestCase):
    def test_pooling_dropdown_defaults_to_legacy_choices(self):
        self.assertEqual(app.pooling_dropdown.value, 'mean')
        self.assertEqual(app.pooling_dropdown.choices, [('mean', 'mean'), ('cls', 'cls')])

    def test_qwen3_pooling_update_forces_last_token(self):
        update = app.update_pooling_for_model(QWEN3_EMBEDDING_MODEL)

        self.assertEqual(update['choices'], ['last-token'])
        self.assertEqual(update['value'], 'last-token')
        self.assertFalse(update['interactive'])

    def test_legacy_pooling_update_restores_mean_cls(self):
        update = app.update_pooling_for_model('bert-base-chinese')

        self.assertEqual(update['choices'], ['mean', 'cls'])
        self.assertEqual(update['value'], 'mean')
        self.assertTrue(update['interactive'])

    def test_qwen3_effective_pooling_ignores_api_pooling_value(self):
        self.assertEqual(app.effective_pooling(QWEN3_EMBEDDING_MODEL, 'mean'), 'last-token')


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