test: topology_raft_disabled: add new util functions

They are shorter and more readable than long CQL queries. We use
them even more in the following commit.

test/topology_raft_disabled/test_raft_upgrade_majority_loss.py

@@ -12,8 +12,8 @@ from test.pylib.manager_client import ManagerClient
 from test.pylib.random_tables import RandomTables
 from test.pylib.util import wait_for_cql_and_get_hosts
 from test.topology.util import reconnect_driver
-from test.topology_raft_disabled.util import restart, enable_raft_and_restart, \
-        wait_until_upgrade_finishes, delete_raft_data, log_run_time
+from test.topology_raft_disabled.util import restart, enable_raft_and_restart, enter_recovery_state, \
+        wait_until_upgrade_finishes, delete_raft_data_and_upgrade_state, log_run_time
 
 
 @pytest.mark.asyncio
@@ -53,7 +53,7 @@ async def test_recovery_after_majority_loss(manager: ManagerClient, random_table
 
     logging.info(f"Entering recovery state on {srv1}")
     host1 = next(h for h in hosts if h.address == srv1.ip_addr)
-    await cql.run_async("update system.scylla_local set value = 'recovery' where key = 'group0_upgrade_state'", host=host1)
+    await enter_recovery_state(cql, host1)
     await restart(manager, srv1)
     cql = await reconnect_driver(manager)
 
@@ -67,8 +67,7 @@ async def test_recovery_after_majority_loss(manager: ManagerClient, random_table
         await manager.remove_node(srv1.server_id, to_remove.server_id, ignore_dead_ips)
 
     logging.info(f"Deleting old Raft data and upgrade state on {host1} and restarting")
-    await delete_raft_data(cql, host1)
-    await cql.run_async("delete from system.scylla_local where key = 'group0_upgrade_state'", host=host1)
+    await delete_raft_data_and_upgrade_state(cql, host1)
     await restart(manager, srv1)
     cql = await reconnect_driver(manager)
 

test/topology_raft_disabled/test_raft_upgrade_stuck.py

@@ -13,8 +13,8 @@ from test.pylib.random_tables import RandomTables
 from test.pylib.rest_client import inject_error_one_shot
 from test.pylib.util import wait_for_cql_and_get_hosts
 from test.topology.util import reconnect_driver
-from test.topology_raft_disabled.util import restart, enable_raft_and_restart, \
-        wait_for_upgrade_state, wait_until_upgrade_finishes, delete_raft_data, log_run_time
+from test.topology_raft_disabled.util import restart, enable_raft_and_restart, enter_recovery_state, \
+        wait_for_upgrade_state, wait_until_upgrade_finishes, delete_raft_data_and_upgrade_state, log_run_time
 
 
 @pytest.mark.asyncio
@@ -56,10 +56,7 @@ async def test_recover_stuck_raft_upgrade(manager: ManagerClient, random_tables:
     # '[shard 0] raft_group0_upgrade - Raft upgrade failed: std::runtime_error (error injection before group 0 upgrade enters synchronize).'
 
     logging.info(f"Setting recovery state on {hosts}")
-    for host in hosts:
-        await cql.run_async(
-                "update system.scylla_local set value = 'recovery' where key = 'group0_upgrade_state'",
-                host=host)
+    await asyncio.gather(*(enter_recovery_state(cql, h) for h in hosts))
 
     logging.info(f"Restarting {others}")
     await asyncio.gather(*(restart(manager, srv) for srv in others))
@@ -85,9 +82,7 @@ async def test_recover_stuck_raft_upgrade(manager: ManagerClient, random_tables:
     await manager.remove_node(others[0].server_id, srv1.server_id)
 
     logging.info(f"Deleting Raft data and upgrade state on {hosts} and restarting")
-    for host in hosts:
-        await delete_raft_data(cql, host)
-        await cql.run_async("delete from system.scylla_local where key = 'group0_upgrade_state'", host=host)
+    await asyncio.gather(*(delete_raft_data_and_upgrade_state(cql, h) for h in hosts))
 
     await asyncio.gather(*(restart(manager, srv) for srv in others))
     cql = await reconnect_driver(manager)

test/topology_raft_disabled/util.py

@@ -77,12 +77,27 @@ async def wait_for_gossip_gen_increase(api: ScyllaRESTAPIClient, gen: int, node_
     logging.info(f"Gossip generation number of {target_ip} is reached {gen} according to {node_ip}")
 
 
+async def enter_recovery_state(cql: Session, host: Host) -> None:
+    await cql.run_async(
+            "update system.scylla_local set value = 'recovery' where key = 'group0_upgrade_state'",
+            host=host)
+
+
 async def delete_raft_data(cql: Session, host: Host) -> None:
     await cql.run_async("truncate table system.discovery", host=host)
     await cql.run_async("truncate table system.group0_history", host=host)
     await cql.run_async("delete value from system.scylla_local where key = 'raft_group0_id'", host=host)
 
 
+async def delete_upgrade_state(cql: Session, host: Host) -> None:
+    await cql.run_async("delete from system.scylla_local where key = 'group0_upgrade_state'", host=host)
+
+
+async def delete_raft_data_and_upgrade_state(cql: Session, host: Host) -> None:
+    await delete_raft_data(cql, host)
+    await delete_upgrade_state(cql, host)
+
+
 def log_run_time(f):
     @functools.wraps(f)
     async def wrapped(*args, **kwargs):