From 34ef8741897cd43be7b17c999f4628389d3bffbb Mon Sep 17 00:00:00 2001
From: Ba Thien Le <bathien.le.pro@gmail.com>
Date: Tue, 27 Aug 2024 12:43:54 +0200
Subject: [PATCH] test: add search with several storages

---
 tests/test_search.py | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/tests/test_search.py b/tests/test_search.py
index 1e732b9..d8835b9 100644
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -44,6 +44,48 @@ def test_search_one_image(client: TestClient) -> None:
     assert len(data["similarities"]) == 1
 
 
+def test_search_one_image_with_storages(client: TestClient) -> None:
+    """
+    Test 'POST /api/search' for one neighbour with several storages.
+
+    Args:
+        client: A test client instance used to send requests to the application.
+    """
+
+    storages = ["test_storage1", "test_storage2"]
+    index_name = "test_index"
+
+    for storage in storages:
+        response = client.post("/api/storages", json={"name": storage})
+        assert response.status_code == 200
+
+        with open("tests/data/image.png", "rb") as file:
+            response = client.post(
+                "/api/images",
+                files={"image": file},
+                params={"storage": storage, "index": index_name},
+            )
+        assert response.status_code == 200
+
+    with open("tests/data/image.png", "rb") as image:
+        response = client.post(
+            "/api/search",
+            files={"image": image},
+            params={
+                "nrt_neigh": "1",
+                "storage": storages,
+                "index": index_name,
+            },
+        )
+
+    data = response.json()
+
+    assert response.status_code == 200
+    assert "similarities" in data
+    assert isinstance(data["similarities"], list)
+    assert len(data["similarities"]) == 1
+
+
 def test_search_images(client: TestClient) -> None:
     """
     Test 'POST /api/search' for several neighbours.