From ef0da14d6fd393702be8bbfc866fe256c82380d0 Mon Sep 17 00:00:00 2001
From: Nadav Har'El <nyh@scylladb.com>
Date: Thu, 6 Oct 2022 20:38:59 +0300
Subject: [PATCH] test/cql-pytest: add simple tests for USE statement

This patch adds a couple of simple tests for the USE statement: that
without USE one cannot create a table without explicitly specifying
a keyspace name, and with USE, it is possible.

Beyond testing these specific feature, this patch also serves as an
example of how to write more tests that need to control the effective USE
setting. Specifically, it adds a "new_cql" function that can be used to
create a new connection with a fresh USE setting. This is necessary
in such tests, because if multiple tests use the same cql fixture
and its single connection, they will share their USE setting and there
is no way to undo or reset it after being set.

Signed-off-by: Nadav Har'El <nyh@scylladb.com>

Closes #11741
---
 test/cql-pytest/test_use.py | 34 ++++++++++++++++++++++++++++++++++
 test/cql-pytest/util.py     | 13 +++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 test/cql-pytest/test_use.py

diff --git a/test/cql-pytest/test_use.py b/test/cql-pytest/test_use.py
new file mode 100644
index 0000000000..0c2b3c444f
--- /dev/null
+++ b/test/cql-pytest/test_use.py
@@ -0,0 +1,34 @@
+# Copyright 2022-present ScyllaDB
+#
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+#############################################################################
+# Tests for the "USE" statement, which modifies the default keyspace used
+# by subsequent statements.
+#
+# Note that because the "USE" statement modifies the state of the current
+# connection, and there is no way to undo its effects (there is no "UNUSE"
+# or way to do an empty "USE"), the following tests should all use a new_cql
+# wrapper over the cql fixture, instead of the cql fixture directly. This
+# wrapper creates a new connection, with its own default USE.
+#############################################################################
+
+import pytest
+from cassandra.protocol import InvalidRequest
+from util import unique_name, new_cql
+
+# Check that CREATE TABLE and DROP TABLE work without an explict keyspace
+# name if a default keyspace name is specified with "USE".
+def test_create_table_use_keyspace(cql, test_keyspace):
+    with new_cql(cql) as ncql:
+        ncql.execute(f'USE {test_keyspace}')
+        table = unique_name()
+        ncql.execute(f'CREATE TABLE {table} (k int PRIMARY KEY)')
+        ncql.execute(f'DROP TABLE {table}')
+
+# Check that without a USE, one cannot CREATE TABLE if the keyspace is not
+# explicitly specified.
+def test_create_table_no_keyspace(cql, test_keyspace):
+    with new_cql(cql) as ncql:
+        with pytest.raises(InvalidRequest, match='No keyspace'):
+            ncql.execute(f'CREATE TABLE {unique_name()} (k int PRIMARY KEY)')
diff --git a/test/cql-pytest/util.py b/test/cql-pytest/util.py
index 69937f264a..557f27ce55 100644
--- a/test/cql-pytest/util.py
+++ b/test/cql-pytest/util.py
@@ -200,6 +200,19 @@ def new_session(cql, username):
         yield session
         session.shutdown()
 
+# new_cql() returns a new object similar to the given cql fixture,
+# connected to the same endpoint but with a separate connection.
+# This can be useful for tests which require a separate connection -
+# for example for testing the "USE" statement (which, after used once
+# on a connection, cannot be undone).
+@contextmanager
+def new_cql(cql):
+    session = cql.cluster.connect()
+    try:
+        yield session
+    finally:
+        session.shutdown()
+
 def project(column_name_string, rows):
     """Returns a list of column values from each of the rows."""
     return [getattr(r, column_name_string) for r in rows]