1  """
 2  Vectors module tests
 3  """
 4  
 5  import os
 6  import tempfile
 7  import unittest
 8  
 9  import numpy as np
10  
11  from txtai.vectors import Vectors, Recovery
12  
13  
14  class TestVectors(unittest.TestCase):
15      """
16      Vectors tests.
17      """
18  
19      def testNotImplemented(self):
20          """
21          Test exceptions for non-implemented methods
22          """
23  
24          vectors = Vectors(None, None, None)
25  
26          self.assertRaises(NotImplementedError, vectors.load, None)
27          self.assertRaises(NotImplementedError, vectors.encode, None)
28  
29      def testNormalize(self):
30          """
31          Test batch normalize and single input normalize are equal
32          """
33  
34          vectors = Vectors(None, None, None)
35  
36          # Generate data
37          data1 = np.random.rand(5, 5).astype(np.float32)
38          data2 = data1.copy()
39  
40          # Keep original data to ensure it changed
41          original = data1.copy()
42  
43          # Normalize data
44          vectors.normalize(data1)
45          for x in data2:
46              vectors.normalize(x)
47  
48          # Test both data arrays are the same and changed from original
49          self.assertTrue(np.allclose(data1, data2))
50          self.assertFalse(np.allclose(data1, original))
51  
52      def testRecovery(self):
53          """
54          Test vectors recovery failure
55          """
56  
57          # Checkpoint directory
58          checkpoint = os.path.join(tempfile.gettempdir(), "recovery")
59          os.makedirs(checkpoint, exist_ok=True)
60  
61          # Create empty file
62          # pylint: disable=R1732
63          f = open(os.path.join(checkpoint, "id"), "w", encoding="utf-8")
64          f.close()
65  
66          # Create the recovery instance with an empty checkpoint file
67          recovery = Recovery(checkpoint, "id", np.load)
68          self.assertIsNone(recovery())