SuperAGI

1
import unittest
2
from superagi.vector_embeddings.pinecone import Pinecone  
3

4

5
class TestPinecone(unittest.TestCase):
6

7
    def setUp(self):
8
        self.uuid = ["id1", "id2"]
9
        self.embeds = ["embed1", "embed2"]
10
        self.metadata = ["metadata1", "metadata2"]
11
        self.pinecone_instance = Pinecone(self.uuid, self.embeds, self.metadata)
12

13
    def test_init(self):
14
        self.assertEqual(self.pinecone_instance.uuid, self.uuid)
15
        self.assertEqual(self.pinecone_instance.embeds, self.embeds)
16
        self.assertEqual(self.pinecone_instance.metadata, self.metadata)
17
    
18
    def test_get_vector_embeddings_from_chunks(self):
19
        expected = {
20
            'vectors': list(zip(self.uuid, self.embeds, self.metadata))
21
        }
22
        result = self.pinecone_instance.get_vector_embeddings_from_chunks()
23
        self.assertEqual(result, expected)
24

25

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