Fix reconfig state storage workload tests (#20994)

7 files changed, 49 insertions, 32 deletions

diff --git a/ydb/tests/stress/reconfig_state_storage_workload/__main__.py b/ydb/tests/stress/reconfig_state_storage_workload/__main__.py
index 0ae23925372..1ea900f7f87 100644
--- a/ydb/tests/stress/reconfig_state_storage_workload/__main__.py
+++ b/ydb/tests/stress/reconfig_state_storage_workload/__main__.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 import argparse
-from ydb.tests.stress.common.common import YdbClient
+import logging
+import sys
 from ydb.tests.stress.reconfig_state_storage_workload.workload import WorkloadRunner
 
 
@@ -8,12 +9,16 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="state storage reconfiguration stability workload", formatter_class=argparse.RawDescriptionHelpFormatter
     )
-    parser.add_argument("--endpoint", default="localhost:2135", help="An endpoint to be used")
-    parser.add_argument("--database", default="Root/test", help="A database to connect")
+    parser.add_argument("--grpc_endpoint", default="localhost:2135", help="An endpoint to be used")
+    parser.add_argument("--http_endpoint", default="http://localhost:8765", help="An endpoint to be used")
+    parser.add_argument("--database", default="/Root", help="A database to connect")
     parser.add_argument("--path", default="olap_workload", help="A path prefix for tables")
+    parser.add_argument("--config_name", default="StateStorage", help="Can be StateStorage / StateStorageBoard / SchemeBoard")
     parser.add_argument("--duration", default=10 ** 9, type=lambda x: int(x), help="A duration of workload in seconds.")
     args = parser.parse_args()
-    client = YdbClient(args.endpoint, args.database, True)
-    client.wait_connection()
-    with WorkloadRunner(client, args.path, args.duration) as runner:
-        runner.run()
+    logger = logging.getLogger("reconfig_state_storage_workload")
+
+    with WorkloadRunner(args.grpc_endpoint, args.http_endpoint, args.database, args.path, args.duration, args.config_name) as runner:
+        if runner.run() < 2:
+            logger.error("Test was not successful.")
+            sys.exit(1)
diff --git a/ydb/tests/stress/reconfig_state_storage_workload/tests/reconfig_state_storage_workload_test.py b/ydb/tests/stress/reconfig_state_storage_workload/tests/reconfig_state_storage_workload_test.py
index d9eb8526811..3d7f8788f05 100644
--- a/ydb/tests/stress/reconfig_state_storage_workload/tests/reconfig_state_storage_workload_test.py
+++ b/ydb/tests/stress/reconfig_state_storage_workload/tests/reconfig_state_storage_workload_test.py
@@ -1,8 +1,10 @@
 # -*- coding: utf-8 -*-
+import os
+import yatest
+
 from ydb.tests.library.harness.kikimr_runner import KiKiMR
 from ydb.tests.library.harness.kikimr_config import KikimrConfigGenerator
 from ydb.tests.library.common.types import Erasure
-from ydb.tests.stress.common.common import YdbClient
 from ydb.tests.library.harness.util import LogLevels
 
 
@@ -30,9 +32,20 @@ class ReconfigStateStorageWorkloadTest(object):
             }
         ))
         cls.cluster.start()
-        cls.client = YdbClient(f'grpc://localhost:{cls.cluster.nodes[1].grpc_port}', '/Root', True)
-        cls.client.wait_connection()
 
     @classmethod
     def teardown_class(cls):
         cls.cluster.stop()
+
+    def do_test(self, config_name):
+        cmd = [
+            yatest.common.binary_path(os.getenv("STRESS_TEST_UTILITY")),
+            "--grpc_endpoint", f"grpc://{self.cluster.nodes[1].host}:{self.cluster.nodes[1].grpc_port}",
+            "--http_endpoint", f"http://{self.cluster.nodes[1].host}:{self.cluster.nodes[1].mon_port}",
+            "--database", "/Root",
+            "--path", "test",
+            "--duration", "120",
+            "--config_name", config_name
+        ]
+
+        yatest.common.execute(cmd, wait=True)
diff --git a/ydb/tests/stress/reconfig_state_storage_workload/tests/test_board_workload.py b/ydb/tests/stress/reconfig_state_storage_workload/tests/test_board_workload.py
index ae62b9c3ea3..3a837a26d3b 100644
--- a/ydb/tests/stress/reconfig_state_storage_workload/tests/test_board_workload.py
+++ b/ydb/tests/stress/reconfig_state_storage_workload/tests/test_board_workload.py
@@ -1,9 +1,7 @@
 # -*- coding: utf-8 -*-
-from ydb.tests.stress.reconfig_state_storage_workload.workload import WorkloadRunner
 from reconfig_state_storage_workload_test import ReconfigStateStorageWorkloadTest
 
 
 class TestReconfigStateStorageBoardWorkload(ReconfigStateStorageWorkloadTest):
     def test_state_storage_board(self):
-        with WorkloadRunner(self.client, self.cluster, 'reconfig_state_storage_board_workload', 120, "StateStorageBoard") as runner:
-            runner.run()
+        self.do_test("StateStorageBoard")
diff --git a/ydb/tests/stress/reconfig_state_storage_workload/tests/test_scheme_board_workload.py b/ydb/tests/stress/reconfig_state_storage_workload/tests/test_scheme_board_workload.py
index 74dbe0da4e0..a6c3c198357 100644
--- a/ydb/tests/stress/reconfig_state_storage_workload/tests/test_scheme_board_workload.py
+++ b/ydb/tests/stress/reconfig_state_storage_workload/tests/test_scheme_board_workload.py
@@ -1,9 +1,7 @@
 # -*- coding: utf-8 -*-
-from ydb.tests.stress.reconfig_state_storage_workload.workload import WorkloadRunner
 from reconfig_state_storage_workload_test import ReconfigStateStorageWorkloadTest
 
 
 class TestReconfigSchemeBoardWorkload(ReconfigStateStorageWorkloadTest):
     def test_scheme_board(self):
-        with WorkloadRunner(self.client, self.cluster, 'reconfig_scheme_board_workload', 120, "SchemeBoard") as runner:
-            runner.run()
+        self.do_test("SchemeBoard")
diff --git a/ydb/tests/stress/reconfig_state_storage_workload/tests/test_state_storage_workload.py b/ydb/tests/stress/reconfig_state_storage_workload/tests/test_state_storage_workload.py
index 285d6e28cf3..ddfd73d123c 100644
--- a/ydb/tests/stress/reconfig_state_storage_workload/tests/test_state_storage_workload.py
+++ b/ydb/tests/stress/reconfig_state_storage_workload/tests/test_state_storage_workload.py
@@ -1,9 +1,7 @@
 # -*- coding: utf-8 -*-
-from ydb.tests.stress.reconfig_state_storage_workload.workload import WorkloadRunner
 from reconfig_state_storage_workload_test import ReconfigStateStorageWorkloadTest
 
 
 class TestReconfigStateStorageWorkload(ReconfigStateStorageWorkloadTest):
     def test_state_storage(self):
-        with WorkloadRunner(self.client, self.cluster, 'reconfig_state_storage_workload', 120, "StateStorage") as runner:
-            runner.run()
+        self.do_test("StateStorage")
diff --git a/ydb/tests/stress/reconfig_state_storage_workload/tests/ya.make b/ydb/tests/stress/reconfig_state_storage_workload/tests/ya.make
index ff84d084c2f..0a39500383b 100644
--- a/ydb/tests/stress/reconfig_state_storage_workload/tests/ya.make
+++ b/ydb/tests/stress/reconfig_state_storage_workload/tests/ya.make
@@ -1,6 +1,6 @@
 PY3TEST()
 INCLUDE(${ARCADIA_ROOT}/ydb/tests/ydbd_dep.inc)
-ENV(YDB_ENABLE_COLUMN_TABLES="true")
+ENV(STRESS_TEST_UTILITY="ydb/tests/stress/reconfig_state_storage_workload/reconfig_state_storage_workload")
 
 TEST_SRCS(
     reconfig_state_storage_workload_test.py
@@ -18,6 +18,7 @@ ENV(IAM_TOKEN="")
 
 DEPENDS(
     ydb/apps/ydb
+    ydb/tests/stress/reconfig_state_storage_workload
 )
 
 PEERDIR(
diff --git a/ydb/tests/stress/reconfig_state_storage_workload/workload/__init__.py b/ydb/tests/stress/reconfig_state_storage_workload/workload/__init__.py
index 9be3a03de88..bec1094c7aa 100644
--- a/ydb/tests/stress/reconfig_state_storage_workload/workload/__init__.py
+++ b/ydb/tests/stress/reconfig_state_storage_workload/workload/__init__.py
@@ -9,6 +9,7 @@ import grpc
 from enum import Enum
 
 from ydb.tests.stress.common.common import WorkloadBase
+from ydb.tests.stress.common.common import YdbClient
 
 logger = logging.getLogger(__name__)
 
@@ -207,10 +208,10 @@ class WorkloadReconfigStateStorage(WorkloadBase):
     loop_cnt = 0
     wait_for = 1
 
-    def __init__(self, client, cluster, prefix, stop, config_name):
+    def __init__(self, client, http_endpoint, prefix, stop, config_name):
         super().__init__(client, prefix, "reconfig_statestorage", stop)
         self.ringGroupActorIdOffset = 1
-        self.cluster = cluster
+        self.http_endpoint = http_endpoint
         self.lock = threading.Lock()
         self.config_name = config_name
 
@@ -219,7 +220,7 @@ class WorkloadReconfigStateStorage(WorkloadBase):
             return f"Reconfig: {self.loop_cnt}"
 
     def do_request(self, json_req):
-        url = f'http://{self.cluster.nodes[1].host}:{self.cluster.nodes[1].mon_port}/actors/nodewarden?page=distconf'
+        url = f'{self.http_endpoint}/actors/nodewarden?page=distconf'
         return requests.post(url, headers={'content-type': 'application/json'}, json=json_req).json()
 
     def do_request_config(self):
@@ -270,9 +271,9 @@ class WorkloadReconfigStateStorage(WorkloadBase):
 
 class WorkloadDiscovery(WorkloadBase):
 
-    def __init__(self, client, cluster, prefix, stop):
+    def __init__(self, client, grpc_endpoint, prefix, stop):
         super().__init__(client, prefix, "discovery", stop)
-        self.cluster = cluster
+        self.grpc_endpoint = grpc_endpoint
         self.lock = threading.Lock()
         self.cnt = 0
 
@@ -281,8 +282,7 @@ class WorkloadDiscovery(WorkloadBase):
             return f"Discovery: {self.cnt}"
 
     def _loop(self):
-        url = "%s:%s" % (self.cluster.nodes[1].host, self.cluster.nodes[1].port)
-        driver_config = ydb.DriverConfig(url, self.client.database)
+        driver_config = ydb.DriverConfig(self.grpc_endpoint, self.client.database)
         while not self.is_stop_requested():
             time.sleep(3)
             resolver = ydb.DiscoveryEndpointsResolver(driver_config)
@@ -301,10 +301,12 @@ class WorkloadDiscovery(WorkloadBase):
 class WorkloadRunner:
     config_name = "StateStorage"
 
-    def __init__(self, client, cluster, path, duration, config_name):
-        self.client = client
+    def __init__(self, grpc_endpoint, http_endpoint, database, path, duration, config_name):
+        self.client = YdbClient(grpc_endpoint, database, True)
+        self.client.wait_connection()
+        self.grpc_endpoint = grpc_endpoint
+        self.http_endpoint = http_endpoint
         self.name = path
-        self.cluster = cluster
         self.tables_prefix = "/".join([self.client.database, self.name])
         self.duration = duration
         self.config_name = config_name
@@ -325,11 +327,12 @@ class WorkloadRunner:
     def run(self):
         stop = threading.Event()
 
+        reconfigWorkload = WorkloadReconfigStateStorage(self.client, self.http_endpoint, self.name, stop, self.config_name)
         workloads = [
             WorkloadTablesCreateDrop(self.client, self.name, stop),
             WorkloadInsertDelete(self.client, self.name, stop),
-            WorkloadReconfigStateStorage(self.client, self.cluster, self.name, stop, self.config_name),
-            WorkloadDiscovery(self.client, self.cluster, self.name, stop)
+            reconfigWorkload,
+            WorkloadDiscovery(self.client, self.grpc_endpoint, self.name, stop)
         ]
         for w in workloads:
             w.start()
@@ -344,3 +347,4 @@ class WorkloadRunner:
         for w in workloads:
             w.join()
         logger.info("Waiting for stop... stopped")
+        return reconfigWorkload.loop_cnt