chore: import upstream snapshot with attribution

test/python/testapi/testcluster.py

@@ -0,0 +1,261 @@
+"""
+Cluster API module tests
+"""
+
+import json
+import os
+import tempfile
+import unittest
+import urllib.parse
+
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from txtai.api import application
+
+# Configuration for an embeddings cluster
+CLUSTER = """
+cluster:
+    shards:
+        - http://127.0.0.1:8002
+        - http://127.0.0.1:8003
+"""
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    """
+    Test HTTP handler.
+    """
+
+    def do_GET(self):
+        """
+        GET request handler.
+        """
+
+        if self.path == "/count":
+            response = 26
+        elif self.path.startswith("/search?query=select"):
+            if "group+by+id" in self.path:
+                response = [{"count(*)": 26}]
+            elif "group+by+text" in self.path:
+                response = [{"count(*)": 12, "text": "This is a test"}, {"count(*)": 14, "text": "And another test"}]
+            elif "group+by+txt" in self.path:
+                response = [{"count(*)": 12, "txt": "This is a test"}, {"count(*)": 14, "txt": "And another test"}]
+            else:
+                if self.server.server_port == 8002:
+                    response = [{"count(*)": 12, "min(indexid)": 0, "max(indexid)": 11, "avg(indexid)": 6.3}]
+                else:
+                    response = [{"count(*)": 16, "min(indexid)": 2, "max(indexid)": 14, "avg(indexid)": 6.7}]
+        elif self.path.startswith("/search"):
+            response = [{"id": 4, "score": 0.40}]
+        else:
+            response = {"result": "ok"}
+
+        # Convert response to string
+        response = json.dumps(response).encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+    def do_POST(self):
+        """
+        POST request handler.
+        """
+
+        if self.path.startswith("/batchsearch"):
+            response = [[{"id": 4, "score": 0.40}], [{"id": 1, "score": 0.40}]]
+        elif self.path.startswith("/delete"):
+            if self.server.server_port == 8002:
+                response = [0]
+            else:
+                response = []
+        else:
+            response = {"result": "ok"}
+
+        response = json.dumps(response).encode("utf-8")
+
+        self.send_response(200)
+        self.send_header("content-type", "application/json")
+        self.send_header("content-length", len(response))
+        self.end_headers()
+
+        self.wfile.write(response)
+        self.wfile.flush()
+
+
+@unittest.skipIf(os.name == "nt", "TestCluster skipped on Windows")
+class TestCluster(unittest.TestCase):
+    """
+    API tests for embeddings clusters
+    """
+
+    @staticmethod
+    @patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
+    def start():
+        """
+        Starts a mock FastAPI client.
+        """
+
+        config = os.path.join(tempfile.gettempdir(), "testapi.yml")
+
+        with open(config, "w", encoding="utf-8") as output:
+            output.write(CLUSTER)
+
+        # Create new application and set on client
+        application.app = application.create()
+        client = TestClient(application.app)
+        application.start()
+
+        return client
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Create API client on creation of class.
+        """
+
+        cls.client = TestCluster.start()
+
+        cls.httpd1 = HTTPServer(("127.0.0.1", 8002), RequestHandler)
+
+        server1 = Thread(target=cls.httpd1.serve_forever, daemon=True)
+        server1.start()
+
+        cls.httpd2 = HTTPServer(("127.0.0.1", 8003), RequestHandler)
+
+        server2 = Thread(target=cls.httpd2.serve_forever, daemon=True)
+        server2.start()
+
+        # Index data
+        cls.client.post("add", json=[{"id": 0, "text": "test"}])
+        cls.client.get("index")
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shutdown mock http server.
+        """
+
+        cls.httpd1.shutdown()
+        cls.httpd2.shutdown()
+
+    def testCount(self):
+        """
+        Test cluster count
+        """
+
+        self.assertEqual(self.client.get("count").json(), 52)
+
+    def testDelete(self):
+        """
+        Test cluster delete
+        """
+
+        self.assertEqual(self.client.post("delete", json=[0]).json(), [0])
+
+    def testDeleteString(self):
+        """
+        Test cluster delete with string id
+        """
+
+        self.assertEqual(self.client.post("delete", json=["0"]).json(), [0])
+
+    def testIds(self):
+        """
+        Test id configurations
+        """
+
+        # String ids
+        self.client.post("add", json=[{"id": "0", "text": "test"}])
+        self.assertEqual(self.client.get("index").status_code, 200)
+
+        # Auto ids
+        self.client.post("add", json=[{"text": "test"}])
+        self.assertEqual(self.client.get("index").status_code, 200)
+
+    def testReindex(self):
+        """
+        Test cluster reindex
+        """
+
+        self.assertEqual(self.client.post("reindex", json={"config": {"path": "sentence-transformers/nli-mpnet-base-v2"}}).status_code, 200)
+
+    def testSearch(self):
+        """
+        Test cluster search
+        """
+
+        # Encode parameters
+        params = json.dumps({"x": 1})
+
+        query = urllib.parse.quote("feel good story")
+        uid = self.client.get(f"search?query={query}&limit=1&weights=0.5&index=default&parameters={params}&graph=False").json()[0]["id"]
+        self.assertEqual(uid, 4)
+
+    def testSearchBatch(self):
+        """
+        Test cluster batch search
+        """
+
+        results = self.client.post(
+            "batchsearch",
+            json={
+                "queries": ["feel good story", "climate change"],
+                "limit": 1,
+                "weights": 0.5,
+                "index": "default",
+                "parameters": [{"x": 1}, {"x": 2}],
+                "graph": False,
+            },
+        ).json()
+
+        uids = [result[0]["id"] for result in results]
+        self.assertEqual(uids, [4, 1])
+
+    def testSQL(self):
+        """
+        Test cluster SQL statement
+        """
+
+        query = urllib.parse.quote("select count(*), min(indexid), max(indexid), avg(indexid) from txtai where text='This is a test'")
+        self.assertEqual(
+            self.client.get(f"search?query={query}").json(), [{"count(*)": 28, "min(indexid)": 0, "max(indexid)": 14, "avg(indexid)": 6.5}]
+        )
+
+        query = urllib.parse.quote("select count(*), text txt from txtai group by txt order by count(*) desc")
+        self.assertEqual(
+            self.client.get(f"search?query={query}").json(),
+            [{"count(*)": 28, "txt": "And another test"}, {"count(*)": 24, "txt": "This is a test"}],
+        )
+
+        query = urllib.parse.quote("select count(*), text from txtai group by text order by count(*) asc")
+        self.assertEqual(
+            self.client.get(f"search?query={query}").json(),
+            [{"count(*)": 24, "text": "This is a test"}, {"count(*)": 28, "text": "And another test"}],
+        )
+
+        query = urllib.parse.quote("select count(*) from txtai group by id order by count(*)")
+        self.assertEqual(self.client.get(f"search?query={query}").json(), [{"count(*)": 52}])
+
+    def testUpsert(self):
+        """
+        Test cluster upsert
+        """
+
+        # Update data
+        self.client.post("add", json=[{"id": 4, "text": "Feel good story: baby panda born"}])
+        self.client.get("upsert")
+
+        # Search for best match
+        query = "feel good story"
+        uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
+
+        self.assertEqual(uid, 4)