chore: Improve integration test for matching engine.

PiperOrigin-RevId: 623575782
googleapis · Apr 10, 2024 · 2d43438 · 2d43438
1 parent aa918e3
commit 2d43438
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/tests/system/aiplatform/test_matching_engine_index.py b/tests/system/aiplatform/test_matching_engine_index.py
@@ -347,6 +347,13 @@ def test_create_get_list_matching_engine_index(self, shared_state):
         # )
         # assert results[0][0].id == 9999
 
+        # FindNeighbors query for public index
+        results = public_index_endpoint.find_neighbors(
+            deployed_index_id=_TEST_DEPLOYED_INDEX_ID_PUBLIC,
+            queries=[_TEST_MATCH_QUERY],
+        )
+        assert results[0][0].id == "0"
+
         # Undeploy index from private endpoint
         my_index_endpoint = my_index_endpoint.undeploy_index(
             deployed_index_id=deployed_index.id