36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
from rag.embedding import EmbeddingError, embed_texts
|
||
|
|
|
||
|
|
|
||
|
|
class EmbeddingTests(unittest.IsolatedAsyncioTestCase):
|
||
|
|
async def test_returns_768_dimension_vectors(self):
|
||
|
|
class Client:
|
||
|
|
async def embed(self, texts):
|
||
|
|
return [[0.1] * 768 for _ in texts]
|
||
|
|
|
||
|
|
vectors = await embed_texts(["基金知识"], client=Client())
|
||
|
|
|
||
|
|
self.assertEqual(len(vectors), 1)
|
||
|
|
self.assertEqual(len(vectors[0]), 768)
|
||
|
|
|
||
|
|
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()
|