Dynamic nodes count for not kube clusters (#14558)

Author: iddqdex

ydb/tests/functional/tpc/lib/conftest.py

@@ -20,7 +20,8 @@ def setup_cluster(cls) -> None:
         YdbCluster.reset(
             ydb_endpoint=f'grpc://{node.host}:{node.grpc_port}',
             ydb_database=f'{cls.cluster.domain_name}/test_db',
-            ydb_mon_port=node.mon_port
+            ydb_mon_port=node.mon_port,
+            dyn_nodes_count=1
         )
         db = f'/{YdbCluster.ydb_database}'
         cls.cluster.create_database(
@@ -29,7 +30,7 @@ def setup_cluster(cls) -> None:
                 'hdd': 1
             }
         )
-        cls.cluster.register_and_start_slots(db, count=1)
+        cls.cluster.register_and_start_slots(db, count=YdbCluster.get_dyn_nodes_count())
         cls.cluster.wait_tenant_up(db)
 
     @classmethod

ydb/tests/olap/lib/ya.make

@@ -11,6 +11,7 @@ PY3_LIBRARY()
     PEERDIR(
         contrib/python/allure-pytest
         contrib/python/allure-python-commons
+        contrib/python/PyYAML
         contrib/python/pytz
         contrib/python/requests
         library/python/testing/yatest_common

ydb/tests/olap/lib/ydb_cluster.py

@@ -3,8 +3,9 @@
 import logging
 import os
 import requests
-from ydb.tests.olap.lib.utils import get_external_param
+import yaml
 import ydb
+from ydb.tests.olap.lib.utils import get_external_param
 from copy import deepcopy
 from time import sleep, time
 from typing import List, Optional
@@ -56,6 +57,7 @@ def __init__(self, desc: dict):
     ydb_mon_port = 8765
     tables_path = get_external_param('tables-path', 'olap_yatests')
     _monitoring_urls: list[YdbCluster.MonitoringUrl] = None
+    _dyn_nodes_count: Optional[int] = None
 
     @classmethod
     def get_monitoring_urls(cls) -> list[YdbCluster.MonitoringUrl]:
@@ -154,10 +156,11 @@ def _create_ydb_driver(endpoint, database, oauth=None, iam_file=None):
             raise
 
     @classmethod
-    def reset(cls, ydb_endpoint, ydb_database, ydb_mon_port):
+    def reset(cls, ydb_endpoint, ydb_database, ydb_mon_port, dyn_nodes_count):
         cls.ydb_endpoint = ydb_endpoint
         cls.ydb_database = ydb_database
         cls.ydb_mon_port = ydb_mon_port
+        cls._dyn_nodes_count = dyn_nodes_count
         cls._ydb_driver = None
 
     @classmethod
@@ -225,6 +228,25 @@ def execute_single_result_query(cls, query, timeout=10):
             LOGGER.error("Cannot connect to YDB")
             raise
 
+    @classmethod
+    def get_dyn_nodes_count(cls) -> int:
+        if cls._dyn_nodes_count is None:
+            cls._dyn_nodes_count = 0
+            if os.getenv('EXPECTED_DYN_NODES_COUNT'):
+                cls._dyn_nodes_count = int(os.getenv('EXPECTED_DYN_NODES_COUNT'))
+            elif os.getenv('CLUSTER_CONFIG'):
+                with open(os.getenv('CLUSTER_CONFIG'), 'r') as r:
+                    yaml_config = yaml.safe_load(r.read())
+                    for domain in yaml_config['domains']:
+                        if domain["domain_name"] == cls.ydb_database:
+                            cls._dyn_nodes_count = domain["domain_name"]["dynamic_slots"]
+                        for db in domain['databases']:
+                            if f'{domain["domain_name"]}/{db["name"]}' == cls.ydb_database:
+                                for cu in db['compute_units']:
+                                    cls._dyn_nodes_count += cu['count']
+
+        return cls._dyn_nodes_count
+
     @classmethod
     @allure.step('Check if YDB alive')
     def check_if_ydb_alive(cls, timeout=10, balanced_paths=None) -> tuple[str, str]:
@@ -241,11 +263,10 @@ def _check_node(n: YdbCluster.Node):
         warnings = []
         try:
             nodes = cls.get_cluster_nodes(db_only=True)
-            expected_nodes_count = os.getenv('EXPECTED_DYN_NODES_COUNT')
+            expected_nodes_count = cls.get_dyn_nodes_count()
             nodes_count = len(nodes)
             if expected_nodes_count:
                 LOGGER.debug(f'Expected nodes count: {expected_nodes_count}')
-                expected_nodes_count = int(expected_nodes_count)
                 if nodes_count < expected_nodes_count:
                     errors.append(f"{expected_nodes_count - nodes_count} nodes from {expected_nodes_count} don't alive")
             ok_node_count = 0

ydb/tests/olap/scenario/conftest.py

@@ -24,6 +24,7 @@ class YdbClusterInstance():
     _temp_ydb_cluster = None
     _endpoint = None
     _database = None
+    _dyn_nodes_count = None
 
     def __init__(self, endpoint, database, config):
         if endpoint is not None:
@@ -38,6 +39,7 @@ def __init__(self, endpoint, database, config):
             self._database = config.domain_name
             self._mon_port = node.mon_port
             self._temp_ydb_cluster = cluster
+            self._dyn_nodes_count = len(cluster.nodes)
         LOGGER.info(f'Using YDB, endpoint:{self._endpoint}, database:{self._database}')
 
     def endpoint(self):
@@ -46,6 +48,9 @@ def endpoint(self):
     def database(self):
         return self._database
 
+    def dyn_nodes_count(self):
+        return self._dyn_nodes_count
+
     def mon_port(self):
         return self._mon_port
 
@@ -65,7 +70,7 @@ def setup_class(cls):
         ydb_endpoint = get_external_param('ydb-endpoint', None)
         ydb_database = get_external_param('ydb-db', "").lstrip('/')
         cls._ydb_instance = YdbClusterInstance(ydb_endpoint, ydb_database, cls._get_cluster_config())
-        YdbCluster.reset(cls._ydb_instance.endpoint(), cls._ydb_instance.database(), cls._ydb_instance.mon_port())
+        YdbCluster.reset(cls._ydb_instance.endpoint(), cls._ydb_instance.database(), cls._ydb_instance.mon_port(), cls._ydb_instance.dyn_nodes_count())
         if not external_param_is_true('reuse-tables'):
             ScenarioTestHelper(None).remove_path(cls.get_suite_name())