import unittest

from rag.embedding import EMBEDDING_DIMENSION, EmbeddingError, embed_texts


class EmbeddingTests(unittest.IsolatedAsyncioTestCase):
    async def test_returns_configured_dimension_vectors(self):
        class Client:
            async def embed(self, texts):
                return [[0.1] * EMBEDDING_DIMENSION for _ in texts]

        vectors = await embed_texts(["基金知识"], client=Client())

        self.assertEqual(len(vectors), 1)
        self.assertEqual(len(vectors[0]), EMBEDDING_DIMENSION)

    async def test_rejects_wrong_embedding_dimension(self):
        class Client:
            async def embed(self, texts):
                return [[0.1] * 3 for _ in texts]

        with self.assertRaises(EmbeddingError):
            await embed_texts(["基金知识"], client=Client())

    async def test_wraps_provider_failure(self):
        class Client:
            async def embed(self, texts):
                raise TimeoutError("embedding timeout")

        with self.assertRaises(EmbeddingError):
            await embed_texts(["基金知识"], client=Client())


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