1  """
  2  Encoding module tests
  3  """
  4  
  5  import base64
  6  import os
  7  import tempfile
  8  import unittest
  9  import urllib.parse
 10  
 11  from io import BytesIO
 12  from unittest.mock import patch
 13  
 14  import msgpack
 15  import numpy as np
 16  import PIL
 17  
 18  from fastapi.testclient import TestClient
 19  
 20  from txtai.api import application
 21  from txtai.api.responses import JSONEncoder
 22  
 23  # pylint: disable=C0411
 24  from utils import Utils
 25  
 26  # Configuration for image storage
 27  INDEX = """
 28  # Allow indexing of documents
 29  writable: %s
 30  
 31  embeddings:
 32      defaults: False
 33      content: True
 34      objects: %s
 35  """
 36  
 37  
 38  class TestEncoding(unittest.TestCase):
 39      """
 40      API tests for response encoding
 41      """
 42  
 43      @staticmethod
 44      @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
 45      def start(yaml):
 46          """
 47          Starts a mock FastAPI client.
 48  
 49          Args:
 50              yaml: input configuration
 51          """
 52  
 53          config = os.path.join(tempfile.gettempdir(), "testapi.yml")
 54  
 55          with open(config, "w", encoding="utf-8") as output:
 56              output.write(yaml)
 57  
 58          # Create new application and set on client
 59          application.app = application.create()
 60          client = TestClient(application.app)
 61          application.start()
 62  
 63          return client
 64  
 65      @classmethod
 66      def setUpClass(cls):
 67          """
 68          Create API client on creation of class.
 69          """
 70  
 71          cls.client = TestEncoding.start(INDEX % ("True", "image"))
 72  
 73      def testImages(self):
 74          """
 75          Test image encoding
 76          """
 77  
 78          with open(Utils.PATH + "/books.jpg", "rb") as f:
 79              self.client.post("addimage", data={"uid": 0}, files={"data": f})
 80              self.client.get("index")
 81  
 82          query = urllib.parse.quote_plus("select id, object from txtai limit 1")
 83          results = self.client.get(f"search?query={query}").json()
 84  
 85          # Test reading image
 86          self.assertIsInstance(PIL.Image.open(BytesIO(base64.b64decode(results[0]["object"]))), PIL.Image.Image)
 87  
 88      def testInvalidInputs(self):
 89          """
 90          Test invalid parameter inputs
 91          """
 92  
 93          response = self.client.post("addimage", data={"uid": [0, 1]}, files={"data": b"123"})
 94          self.assertEqual(response.status_code, 422)
 95  
 96          response = self.client.post("addobject", data={"uid": [0, 1]}, files={"data": b"123"})
 97          self.assertEqual(response.status_code, 422)
 98  
 99      def testInvalidJSON(self):
100          """
101          Test that invalid JSON raises an exception
102          """
103  
104          with self.assertRaises(TypeError):
105              JSONEncoder().encode(np.random.rand(1, 1))
106  
107      def testMessagePack(self):
108          """
109          Test message pack encoding
110          """
111  
112          # Validate binary encoding
113          results = self.client.get("count", headers={"Accept": "application/msgpack"}).content
114          self.assertEqual(results, b"\x01")
115  
116          # Validate query result
117          query = urllib.parse.quote_plus("select id, object from txtai limit 1")
118          results = self.client.get(f"search?query={query}", headers={"Accept": "application/msgpack"}).content
119          results = msgpack.unpackb(results)
120  
121          # Test reading image
122          self.assertIsInstance(PIL.Image.open(BytesIO(results[0]["object"])), PIL.Image.Image)
123  
124      def testObjects(self):
125          """
126          Test object encoding
127          """
128  
129          # Recreate model with standard object encoding
130          self.client = TestEncoding.start(INDEX % ("True", "True"))
131  
132          # Test various formats
133          self.client.post("addobject", data={"uid": "id0"}, files={"data": b"1234"})
134          self.client.post("addobject", files={"data": b"ABC"})
135          self.client.post("addobject", data={"uid": "id1", "field": "object"}, files={"data": b"A1234"})
136          self.client.get("index")
137  
138          query = urllib.parse.quote_plus("select id, object from txtai where id = 'id0' limit 1")
139          results = self.client.get(f"search?query={query}").json()
140          self.assertEqual(base64.b64decode(results[0]["object"]), b"1234")
141  
142          # Test with messagepack encoding
143          results = self.client.get(f"search?query={query}", headers={"Accept": "application/msgpack"}).content
144          results = msgpack.unpackb(results)
145          self.assertEqual(results[0]["object"], b"1234")
146  
147          count = self.client.get("count").json()
148          self.assertEqual(count, 3)
149  
150      def testReadOnly(self):
151          """
152          Test read only indexes
153          """
154  
155          # Recreate model with standard object encoding
156          self.client = TestEncoding.start(INDEX % ("False", "True"))
157  
158          # Test errors raised for write operations
159          with open(Utils.PATH + "/books.jpg", "rb") as f:
160              response = self.client.post("addimage", data={"uid": 0}, files={"data": f})
161              self.assertEqual(response.status_code, 403)
162  
163          self.assertEqual(self.client.post("addobject", data={"uid": 0}, files={"data": b"1234"}).status_code, 403)