SuperAGI

1
import unittest
2
from superagi.vector_embeddings.base import VectorEmbeddings
3
from superagi.vector_embeddings.weaviate import Weaviate
4

5
class TestWeaviate(unittest.TestCase):
6

7
    def setUp(self):
8
        self.weaviate = Weaviate(uuid="1234", embeds=[0.1, 0.2, 0.3, 0.4], metadata={"info": "sample data"})
9

10
    def test_init(self):
11
        self.assertEqual(self.weaviate.uuid, "1234")
12
        self.assertEqual(self.weaviate.embeds, [0.1, 0.2, 0.3, 0.4])
13
        self.assertEqual(self.weaviate.metadata, {"info": "sample data"})
14

15
    def test_get_vector_embeddings_from_chunks(self):
16
        expected_result = {
17
            "ids": "1234",
18
            "data_object": {"info": "sample data"},
19
            "vectors": [0.1, 0.2, 0.3, 0.4]
20
        }
21
        self.assertEqual(self.weaviate.get_vector_embeddings_from_chunks(), expected_result)
22

23

24
if __name__ == '__main__':
25
    unittest.main()
26