fix(manager): extend some load/stress methods to handle keyspace names

Manager's _generate_load method has been worked to take keyspace_name
as a parameter instead of keyspace_name_to_replace. That allows to
generate the load for specific keyspace and makes the code more clear
and understandable.

assemble_and_run_all_stress_cmd_by_ks_names method was introduced
to run stress commands for keyspaces with provided names. It was
decided to keep it as separate method to not overcomplicate the
existing one (assemble_and_run_all_stress_cmd).

mgmt_cli_test.py

@@ -200,30 +200,31 @@ def restore_backup_with_task(self, mgr_cluster, snapshot_tag, timeout, restore_s
         if restore_schema:
             self.db_cluster.restart_scylla()  # After schema restoration, you should restart the nodes
 
-    def run_verification_read_stress(self):
+    def run_verification_read_stress(self, ks_names=None):
         stress_queue = []
         stress_cmd = self.params.get('stress_read_cmd')
         keyspace_num = self.params.get('keyspace_num')
         InfoEvent(message='Starting read stress for data verification').publish()
         stress_start_time = datetime.now()
-        self.assemble_and_run_all_stress_cmd(stress_queue, stress_cmd, keyspace_num)
+        if ks_names:
+            self.assemble_and_run_all_stress_cmd_by_ks_names(stress_queue, stress_cmd, ks_names)
+        else:
+            self.assemble_and_run_all_stress_cmd(stress_queue, stress_cmd, keyspace_num)
         for stress in stress_queue:
             self.verify_stress_thread(cs_thread_pool=stress)
         stress_run_time = datetime.now() - stress_start_time
         InfoEvent(message=f'The read stress run was completed. Total run time: {stress_run_time}').publish()
 
-    def _generate_load(self, keyspace_name_to_replace=None):
+    def _generate_load(self, keyspace_name: str = None):
         self.log.info('Starting c-s write workload')
         stress_cmd = self.params.get('stress_cmd')
-        if keyspace_name_to_replace:
-            stress_cmd = stress_cmd.replace("keyspace1", keyspace_name_to_replace)
-        stress_thread = self.run_stress_thread(stress_cmd=stress_cmd)
+        stress_thread = self.run_stress_thread(stress_cmd=stress_cmd, keyspace_name=keyspace_name)
         self.log.info('Sleeping for 15s to let cassandra-stress run...')
         time.sleep(15)
         return stress_thread
 
-    def generate_load_and_wait_for_results(self, keyspace_name_to_replace=None):
-        load_thread = self._generate_load(keyspace_name_to_replace=keyspace_name_to_replace)
+    def generate_load_and_wait_for_results(self, keyspace_name: str = None):
+        load_thread = self._generate_load(keyspace_name=keyspace_name)
         load_results = load_thread.get_results()
         self.log.info(f'load={load_results}')
 

mgmt_upgrade_test.py

@@ -106,7 +106,7 @@ def test_upgrade(self):  # pylint: disable=too-many-locals,too-many-statements
                 f"Unknown failure in task {rerunning_backup_task.id}"
 
         with self.subTest("Creating a backup task and stopping it"):
-            self.generate_load_and_wait_for_results(keyspace_name_to_replace="keyspace2")
+            self.generate_load_and_wait_for_results(keyspace_name="keyspace2")
             legacy_args = "--force" if manager_tool.sctool.client_version.startswith("2.1") else None
             pausable_backup_task = mgr_cluster.create_backup_task(
                 cron=create_cron_list_from_timedelta(minutes=1),

sdcm/utils/loader_utils.py

@@ -65,21 +65,29 @@ def _create_counter_table(self):
 
     def assemble_and_run_all_stress_cmd(self, stress_queue, stress_cmd, keyspace_num):
         if stress_cmd:
+            round_robin = self.params.get('round_robin')
             # Stress: Same as in prepare_write - allow the load to be spread across all loaders when using multi ks
-            if keyspace_num > 1 and self.params.get('round_robin'):
+            if keyspace_num > 1 and round_robin:
                 self.log.debug("Using round_robin for multiple Keyspaces...")
                 for i in range(1, keyspace_num + 1):
                     keyspace_name = self._get_keyspace_name(i)
-                    params = {'keyspace_name': keyspace_name, 'round_robin': True, 'stress_cmd': stress_cmd}
-
+                    params = {'keyspace_name': keyspace_name, 'stress_cmd': stress_cmd, 'round_robin': round_robin}
                     self._run_all_stress_cmds(stress_queue, params)
 
             # The old method when we run all stress_cmds for all keyspace on the same loader, or in round-robin if defined in test yaml
             else:
-                params = {'keyspace_num': keyspace_num, 'stress_cmd': stress_cmd,
-                          'round_robin': self.params.get('round_robin')}
+                params = {'keyspace_num': keyspace_num, 'stress_cmd': stress_cmd, 'round_robin': round_robin}
                 self._run_all_stress_cmds(stress_queue, params)
 
+    def assemble_and_run_all_stress_cmd_by_ks_names(self, stress_queue, stress_cmd, ks_names=None):
+        for ks_name in ks_names:
+            params = {
+                'keyspace_name': ks_name,
+                'stress_cmd': stress_cmd,
+                'round_robin': self.params.get('round_robin')
+            }
+            self._run_all_stress_cmds(stress_queue, params)
+
     def _run_all_stress_cmds(self, stress_queue, params):
         stress_cmds = params['stress_cmd']
         if not isinstance(stress_cmds, list):

test-cases/upgrades/manager-upgrade.yaml

@@ -1,7 +1,7 @@
 test_duration: 360
 
-stress_cmd: "cassandra-stress write cl=QUORUM n=1200300 -schema 'replication(strategy=NetworkTopologyStrategy,replication_factor=3) keyspace=keyspace1' -mode cql3 native -rate threads=200 -pop seq=400200300..600200300"
-stress_read_cmd: "cassandra-stress read cl=QUORUM n=1200300 -schema 'replication(strategy=NetworkTopologyStrategy,replication_factor=3) keyspace=keyspace1' -mode cql3 native -rate threads=200 -pop seq=400200300..600200300"
+stress_cmd: "cassandra-stress write cl=QUORUM n=1200300 -schema 'replication(strategy=NetworkTopologyStrategy,replication_factor=3)' -mode cql3 native -rate threads=200 -pop seq=400200300..600200300"
+stress_read_cmd: "cassandra-stress read cl=QUORUM n=1200300 -schema 'replication(strategy=NetworkTopologyStrategy,replication_factor=3)' -mode cql3 native -rate threads=200 -pop seq=400200300..600200300"
 
 instance_type_db: 'i4i.large'
 instance_type_loader: 'c6i.large'