From 8ea862f1e8fb53a95773e778cd5b2fa9fe41d6a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Hudobski?= <michal.hudobski@scylladb.com>
Date: Mon, 12 May 2025 10:48:17 +0200
Subject: [PATCH] test/cqlpy: add custom index tests

Unit tests checking the behavior
of the added support for create
custom index statement
---
 test/cqlpy/test_vector_index.py | 77 +++++++++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)
 create mode 100644 test/cqlpy/test_vector_index.py

diff --git a/test/cqlpy/test_vector_index.py b/test/cqlpy/test_vector_index.py
new file mode 100644
index 0000000000..6ba09bf19a
--- /dev/null
+++ b/test/cqlpy/test_vector_index.py
@@ -0,0 +1,77 @@
+# Copyright 2025-present ScyllaDB
+#
+# SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
+
+###############################################################################
+# Tests for vector indexes
+###############################################################################
+
+import pytest
+from .util import new_test_table, is_scylla
+from cassandra.protocol import InvalidRequest, ConfigurationException
+
+
+def test_create_vector_search_index(cql, test_keyspace, scylla_only):
+    schema = 'p int primary key, v vector<float, 3>'
+    with new_test_table(cql, test_keyspace, schema) as table:
+        cql.execute(f"CREATE CUSTOM INDEX ON {table}(v) USING 'vector_index'")
+
+
+
+def test_create_vector_search_index_without_custom_keyword(cql, test_keyspace):
+    schema = 'p int primary key, v vector<float, 3>'
+    with new_test_table(cql, test_keyspace, schema) as table:
+        if is_scylla(cql):
+            custom_class = 'vector_index'
+        else:
+            custom_class =  'sai'
+        
+        cql.execute(f"CREATE INDEX ON {table}(v) USING '{custom_class}'")
+
+def test_create_custom_index_with_invalid_class(cql, test_keyspace):
+    schema = 'p int primary key, v vector<float, 3>'
+    with new_test_table(cql, test_keyspace, schema) as table:
+        invalid_custom_class = "invalid.custom.class"
+        with pytest.raises((InvalidRequest, ConfigurationException), match=r"Non-supported custom class|Unable to find"):
+            cql.execute(f"CREATE CUSTOM INDEX ON {table}(v) USING '{invalid_custom_class}'")
+
+def test_create_custom_index_without_custom_class(cql, test_keyspace):
+    schema = 'p int primary key, v vector<float, 3>'
+    with new_test_table(cql, test_keyspace, schema) as table:
+        with pytest.raises((InvalidRequest, ConfigurationException), match=r"CUSTOM index requires specifying|Unable to find"):
+            cql.execute(f"CREATE CUSTOM INDEX ON {table}(v)")
+
+@pytest.mark.xfail(reason="Scylla doesn't validate vector indexes, as they are not implemented yet.")
+def test_create_vector_search_index_on_nonvector_column(cql, test_keyspace, scylla_only):
+    schema = 'p int primary key, v int'
+    with new_test_table(cql, test_keyspace, schema) as table:
+        with pytest.raises(InvalidRequest):
+            cql.execute(f"CREATE CUSTOM INDEX ON {table}(v) USING 'vector_index'")
+        
+
+def test_describe_custom_index(cql, test_keyspace):
+    schema = 'p int primary key, v1 vector<float, 3>, v2 vector<float, 3>'
+    with new_test_table(cql, test_keyspace, schema) as table:
+        # Cassandra inserts a space between the table name and parentheses,
+        # Scylla doesn't. This difference doesn't matter because both are
+        # valid CQL commands
+        # Scylla doesn't support sai custom class.
+        if is_scylla(cql):
+            maybe_space = ''
+            custom_class = 'vector_index'
+        else:
+            maybe_space = ' '
+            custom_class =  'sai'
+
+
+        create_idx_a = f"CREATE INDEX custom ON {table}(v1) USING '{custom_class}'"
+        create_idx_b = f"CREATE CUSTOM INDEX custom1 ON {table}(v2) USING '{custom_class}'"
+
+        cql.execute(create_idx_a)
+        cql.execute(create_idx_b)
+
+        a_desc = cql.execute(f"DESC INDEX {test_keyspace}.custom").one().create_statement
+        b_desc = cql.execute(f"DESC INDEX {test_keyspace}.custom1").one().create_statement
+
+        assert f"CREATE CUSTOM INDEX custom ON {table}{maybe_space}(v1) USING '{custom_class}'" in a_desc
+        assert f"CREATE CUSTOM INDEX custom1 ON {table}{maybe_space}(v2) USING '{custom_class}'" in b_desc
\ No newline at end of file