mirrors/scylladb
1
0
Fork 0
mirror of https://github.com/scylladb/scylladb.git synced 2026-04-20 16:40:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

vector_search: test: refactor boilerplate setup

Browse Source 
The test boilerplate setup for some vector store client tests
has been extracted to a common function.
This commit is contained in:
Karol Nowacki
2026-03-11 20:26:23 +01:00
parent 6bc88e817f
commit a32e4bb9f4
1 changed files with 41 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

102
test/vector_search/vector_store_client_test.cc
View File

@@ -83,6 +83,25 @@ timeout_config make_query_timeout(std::chrono::seconds timeout) {
return cfg;
}
future<> do_with_vector_store_mock(std::function<future<>(cql_test_env&, vs_mock_server&)> func) {
auto server = co_await make_vs_mock_server();
auto cfg = make_config();
cfg.db_config->vector_store_primary_uri.set(format("http://server.node:{}", server->port()));
co_await do_with_cql_env(
[&](cql_test_env& env) -> future<> {
co_await create_test_table(env, "ks", "test");
auto& vs = env.local_qp().vector_store_client();
configure(vs).with_dns({{"server.node", std::vector<std::string>{server->host()}}});
vs.start_background_tasks();
co_await func(env, *server);
},
cfg)
.finally(seastar::coroutine::lambda([&] -> future<> {
co_await server->stop();
}));
}
} // namespace
BOOST_AUTO_TEST_CASE(vector_store_client_test_ctor) {
@@ -1222,32 +1241,19 @@ SEASTAR_TEST_CASE(vector_store_client_abort_due_to_query_timeout) {
/// Verify that the HTTP error description from the vector store is propagated
/// through the CQL interface as part of the invalid_request_exception message.
SEASTAR_TEST_CASE(vector_store_client_cql_error_contains_http_error_description) {
auto server = co_await make_vs_mock_server();
co_await do_with_vector_store_mock([](cql_test_env& env, vs_mock_server& server) -> future<> {
co_await env.execute_cql("CREATE CUSTOM INDEX idx ON ks.test (embedding) USING 'vector_index'");
auto cfg = make_config();
cfg.db_config->vector_store_primary_uri.set(format("http://server.node:{}", server->port()));
co_await do_with_cql_env(
[&](cql_test_env& env) -> future<> {
auto schema = co_await create_test_table(env, "ks", "test");
auto& vs = env.local_qp().vector_store_client();
configure(vs).with_dns({{"server.node", std::vector<std::string>{server->host()}}});
vs.start_background_tasks();
co_await env.execute_cql("CREATE CUSTOM INDEX idx ON ks.test (embedding) USING 'vector_index'");
// Configure mock to return 404 with a specific error message
server.next_ann_response({status_type::not_found, "index does not exist"});
// Configure mock to return 404 with a specific error message
server->next_ann_response({status_type::not_found, "index does not exist"});
BOOST_CHECK_EXCEPTION(co_await env.execute_cql("SELECT * FROM ks.test ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 5;"),
exceptions::invalid_request_exception, [](const exceptions::invalid_request_exception& ex) {
auto msg = std::string(ex.what());
// Verify the error message contains both the HTTP status and the error description
return msg.find("404") != std::string::npos && msg.find("index does not exist") != std::string::npos;
});
},
cfg)
.finally(seastar::coroutine::lambda([&] -> future<> {
co_await server->stop();
}));
BOOST_CHECK_EXCEPTION(co_await env.execute_cql("SELECT * FROM ks.test ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 5;"),
exceptions::invalid_request_exception, [](const exceptions::invalid_request_exception& ex) {
auto msg = std::string(ex.what());
// Verify the error message contains both the HTTP status and the error description
return msg.find("404") != std::string::npos && msg.find("index does not exist") != std::string::npos;
});
});
}
// Create a vector index with an additional filtering column.
@@ -1257,46 +1263,20 @@ SEASTAR_TEST_CASE(vector_store_client_cql_error_contains_http_error_description)
// ANN ordering by vector requires the column to be indexed using 'vector_index'.
// Reproduces SCYLLADB-635.
SEASTAR_TEST_CASE(vector_store_client_vector_index_with_additional_filtering_column) {
auto server = co_await make_vs_mock_server();
co_await do_with_vector_store_mock([](cql_test_env& env, vs_mock_server&) -> future<> {
// Create a vector index on the embedding column, including ck1 for filtered ANN search support.
co_await env.execute_cql("CREATE CUSTOM INDEX idx ON ks.test (embedding, ck1) USING 'vector_index'");
auto cfg = make_config();
cfg.db_config->vector_store_primary_uri.set(format("http://server.node:{}", server->port()));
co_await do_with_cql_env(
[&](cql_test_env& env) -> future<> {
auto schema = co_await create_test_table(env, "ks", "test");
auto& vs = env.local_qp().vector_store_client();
configure(vs).with_dns({{"server.node", std::vector<std::string>{server->host()}}});
vs.start_background_tasks();
// Create a vector index on the embedding column, including ck1 for filtered ANN search support.
auto result = co_await env.execute_cql("CREATE CUSTOM INDEX idx ON ks.test (embedding, ck1) USING 'vector_index'");
BOOST_CHECK_NO_THROW(co_await env.execute_cql("SELECT * FROM ks.test ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 5;"));
},
cfg)
.finally(seastar::coroutine::lambda([&] -> future<> {
co_await server->stop();
}));
BOOST_CHECK_NO_THROW(co_await env.execute_cql("SELECT * FROM ks.test ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 5;"));
});
}
SEASTAR_TEST_CASE(vector_store_client_local_vector_index) {
auto server = co_await make_vs_mock_server();
co_await do_with_vector_store_mock([](cql_test_env& env, vs_mock_server&) -> future<> {
// Create a local vector index on the 'embedding' column.
co_await env.execute_cql("CREATE CUSTOM INDEX idx ON ks.test ((pk1, pk2), embedding) USING 'vector_index'");
auto cfg = make_config();
cfg.db_config->vector_store_primary_uri.set(format("http://server.node:{}", server->port()));
co_await do_with_cql_env(
[&](cql_test_env& env) -> future<> {
auto schema = co_await create_test_table(env, "ks", "test");
auto& vs = env.local_qp().vector_store_client();
configure(vs).with_dns({{"server.node", std::vector<std::string>{server->host()}}});
vs.start_background_tasks();
// Create a local vector index on the 'embedding' column.
auto result = co_await env.execute_cql("CREATE CUSTOM INDEX idx ON ks.test ((pk1, pk2), embedding) USING 'vector_index'");
BOOST_CHECK_NO_THROW(
co_await env.execute_cql("SELECT * FROM ks.test WHERE pk1 = 1 AND pk2 = 2 ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 5;"));
},
cfg)
.finally(seastar::coroutine::lambda([&] -> future<> {
co_await server->stop();
}));
BOOST_CHECK_NO_THROW(
co_await env.execute_cql("SELECT * FROM ks.test WHERE pk1 = 1 AND pk2 = 2 ORDER BY embedding ANN OF [0.1, 0.2, 0.3] LIMIT 5;"));
});
}