1  """
 2  Client module tests
 3  """
 4  
 5  import os
 6  import time
 7  import tempfile
 8  
 9  from txtai.embeddings import Embeddings
10  
11  from .testrdbms import Common
12  
13  
14  # pylint: disable=R0904
15  class TestClient(Common.TestRDBMS):
16      """
17      Embeddings with content stored in a client RDBMS.
18      """
19  
20      @classmethod
21      def setUpClass(cls):
22          """
23          Initialize test data.
24          """
25  
26          cls.data = [
27              "US tops 5 million confirmed virus cases",
28              "Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
29              "Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
30              "The National Park Service warns against sacrificing slower friends in a bear attack",
31              "Maine man wins $1M from $25 lottery ticket",
32              "Make huge profits without work, earn up to $100,000 a day",
33          ]
34  
35          # Content backend
36          cls.backend = None
37  
38          # Create embeddings model, backed by sentence-transformers & transformers
39          cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
40  
41      @classmethod
42      def tearDownClass(cls):
43          """
44          Cleanup data.
45          """
46  
47          if cls.embeddings:
48              cls.embeddings.close()
49  
50      def setUp(self):
51          """
52          Set unique database path for each test.
53          """
54  
55          # Generate unique database path and set on embeddings
56          path = os.path.join(tempfile.gettempdir(), f"{int(time.time() * 1000)}.sqlite")
57          self.backend = f"sqlite:///{path}"
58  
59          self.embeddings.config["content"] = self.backend
60  
61      def testSchema(self):
62          """
63          Test database creation with a specified schema
64          """
65  
66          # Default sequence id
67          embeddings = Embeddings(path="sentence-transformers/nli-mpnet-base-v2", content=self.backend, schema="txtai")
68          embeddings.index(self.data)
69  
70          result = embeddings.search("feel good story", 1)[0]
71          self.assertEqual(result["text"], self.data[4])