Wait until db is populated in connector/fq tests (#10047)

yumkam · web-flow · commit 912ae27fb96d · 2024-10-09T16:32:08.000+03:00
diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/init/01_basic.sh b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/init/01_basic.sh
@@ -162,6 +162,10 @@ set -ex
       (2, @@{ "TODO": "unicode" }@@, CAST(@@{ "TODO": "unicode" }@@AS Json)),
       (3, @@{ }@@, NULL);
     COMMIT;
+    CREATE TABLE dummy_table (name String, cnt Uint64, PRIMARY KEY(name));
+    COMMIT;
+    INSERT INTO dummy_table (name, cnt) SELECT "json", COUNT(*) FROM json;
+    COMMIT;
   '
 
 echo $(date +"%T.%6N") "SUCCESS"
diff --git a/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/test.py b/ydb/library/yql/providers/generic/connector/tests/datasource/ydb/test.py
@@ -1,14 +1,7 @@
 import pytest
-from datetime import datetime
-import time
-from typing import Sequence
-
-import yatest.common
 
 from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind
 from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings
-from ydb.library.yql.providers.generic.connector.tests.utils.log import make_logger
-from ydb.library.yql.providers.generic.connector.tests.utils.docker_compose import DockerComposeHelper
 from ydb.library.yql.providers.generic.connector.tests.utils.run.runners import runner_types, configure_runner
 import ydb.library.yql.providers.generic.connector.tests.utils.scenario.ydb as scenario
 import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_positive_common as select_positive_common
@@ -19,41 +12,8 @@
 from conftest import docker_compose_dir
 from collection import Collection
 
-LOGGER = make_logger(__name__)
-
-
-class OneTimeWaiter:
-    __launched: bool = False
-
-    def __init__(self, expected_tables: Sequence[str]):
-        docker_compose_file_relative_path = str(docker_compose_dir / 'docker-compose.yml')
-        docker_compose_file_abs_path = yatest.common.source_path(docker_compose_file_relative_path)
-        self.docker_compose_helper = DockerComposeHelper(docker_compose_yml_path=docker_compose_file_abs_path)
-        self.expected_tables = set(expected_tables)
-
-    def wait(self):
-        if self.__launched:
-            return
-
-        # This should be enough for tables to initialize
-        start = datetime.now()
-
-        timeout = 600
-        while (datetime.now() - start).total_seconds() < timeout:
-            self.actual_tables = set(self.docker_compose_helper.list_ydb_tables())
-
-            # check if all the required tables have been created
-            if self.expected_tables <= self.actual_tables:
-                self.__launched = True
-                return
-
-            LOGGER.warning(f"Not enough YDB tables: expected={self.expected_tables}, actual={self.actual_tables}")
-            time.sleep(5)
-
-        raise ValueError(f"YDB was not able to initialize in {timeout} seconds, latest table set: {self.actual_tables}")
-
-
-one_time_waiter = OneTimeWaiter(
+one_time_waiter = scenario.OneTimeWaiter(
+    docker_compose_file_path=str(docker_compose_dir / 'docker-compose.yml'),
     expected_tables=[
         "column_selection_A_b_C_d_E",
         "column_selection_COL1",
@@ -68,6 +28,7 @@ def wait(self):
         "pushdown",
         "unsupported_types",
         "json",
+        "dummy_table",
     ]
 )
 
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/scenario/ya.make b/ydb/library/yql/providers/generic/connector/tests/utils/scenario/ya.make
@@ -20,6 +20,7 @@ PEERDIR(
     ydb/library/yql/providers/generic/connector/tests/utils
     ydb/library/yql/providers/generic/connector/tests/utils/clients
     ydb/library/yql/providers/generic/connector/tests/utils/run
+    ydb/library/yql/providers/generic/connector/tests/common_test_cases
 )
 
 END()
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/scenario/ydb.py b/ydb/library/yql/providers/generic/connector/tests/utils/scenario/ydb.py
@@ -1,3 +1,11 @@
+from datetime import datetime
+import time
+from typing import Sequence
+
+import yatest.common
+from ydb.library.yql.providers.generic.connector.tests.utils.log import make_logger
+from ydb.library.yql.providers.generic.connector.tests.utils.docker_compose import DockerComposeHelper
+
 from ydb.library.yql.providers.generic.connector.tests.utils.comparator import assert_data_outs_equal
 from ydb.library.yql.providers.generic.connector.tests.utils.settings import Settings
 from ydb.library.yql.providers.generic.connector.tests.utils.run.parent import Runner
@@ -7,6 +15,38 @@
 
 # import ydb.library.yql.providers.generic.connector.tests.common_test_cases.select_missing_database as tc_select_missing_database
 
+LOGGER = make_logger(__name__)
+
+
+class OneTimeWaiter:
+    __launched: bool = False
+
+    def __init__(self, docker_compose_file_path: str, expected_tables: Sequence[str]):
+        docker_compose_file_abs_path = yatest.common.source_path(docker_compose_file_path)
+        self.docker_compose_helper = DockerComposeHelper(docker_compose_yml_path=docker_compose_file_abs_path)
+        self.expected_tables = set(expected_tables)
+
+    def wait(self):
+        if self.__launched:
+            return
+
+        # This should be enough for tables to initialize
+        start = datetime.now()
+
+        timeout = 600
+        while (datetime.now() - start).total_seconds() < timeout:
+            self.actual_tables = set(self.docker_compose_helper.list_ydb_tables())
+
+            # check if all the required tables have been created
+            if self.expected_tables <= self.actual_tables:
+                self.__launched = True
+                return
+
+            LOGGER.warning(f"Not enough YDB tables: expected={self.expected_tables}, actual={self.actual_tables}")
+            time.sleep(5)
+
+        raise ValueError(f"YDB was not able to initialize in {timeout} seconds, latest table set: {self.actual_tables}")
+
 
 def select_positive(
     test_name: str,
diff --git a/ydb/tests/fq/generic/analytics/conftest.py b/ydb/tests/fq/generic/analytics/conftest.py
@@ -11,10 +11,14 @@
 
 from ydb.tests.fq.generic.utils.settings import Settings
 
+from typing import Final
+
+docker_compose_file_path: Final = "ydb/tests/fq/generic/analytics/docker-compose.yml"
+
 
 @pytest.fixture
 def settings() -> Settings:
-    return Settings.from_env(docker_compose_file_path='ydb/tests/fq/generic/analytics/docker-compose.yml')
+    return Settings.from_env(docker_compose_file_path=docker_compose_file_path)
 
 
 @pytest.fixture
diff --git a/ydb/tests/fq/generic/analytics/test_join.py b/ydb/tests/fq/generic/analytics/test_join.py
@@ -6,6 +6,14 @@
 
 from ydb.tests.tools.fq_runner.fq_client import FederatedQueryClient
 from ydb.tests.fq.generic.utils.settings import Settings
+from ydb.library.yql.providers.generic.connector.tests.utils.scenario.ydb import OneTimeWaiter
+import conftest
+
+
+one_time_waiter = OneTimeWaiter(
+    docker_compose_file_path=conftest.docker_compose_file_path,
+    expected_tables=["join_table", "dummy_table"],
+)
 
 
 class TestJoinAnalytics:
@@ -47,6 +55,8 @@ def test_simple(self, fq_client: FederatedQueryClient, settings: Settings, query
             password=settings.postgresql.password,
         )
 
+        one_time_waiter.wait()
+
         fq_client.create_ydb_connection(
             name=ydb_conn_name,
             database_id=settings.ydb.dbname,
diff --git a/ydb/tests/fq/generic/analytics/test_ydb.py b/ydb/tests/fq/generic/analytics/test_ydb.py
@@ -7,6 +7,14 @@
 
 from ydb.tests.tools.fq_runner.fq_client import FederatedQueryClient
 from ydb.tests.fq.generic.utils.settings import Settings
+from ydb.library.yql.providers.generic.connector.tests.utils.scenario.ydb import OneTimeWaiter
+import conftest
+
+
+one_time_waiter = OneTimeWaiter(
+    docker_compose_file_path=conftest.docker_compose_file_path,
+    expected_tables=["simple_table", "dummy_table"],
+)
 
 
 class TestYdb:
@@ -20,6 +28,8 @@ def test_simple(self, fq_client: FederatedQueryClient, settings: Settings):
         conn_name = f'conn_{table_name}'
         query_name = f'query_{table_name}'
 
+        one_time_waiter.wait()
+
         fq_client.create_ydb_connection(
             name=conn_name,
             database_id=settings.ydb.dbname,
diff --git a/ydb/tests/fq/generic/analytics/ya.make b/ydb/tests/fq/generic/analytics/ya.make
@@ -52,6 +52,7 @@ IF (OPENSOURCE)
 ENDIF()
 
 PEERDIR(
+    ydb/library/yql/providers/generic/connector/tests/utils/scenario
     ydb/tests/fq/generic/utils
     library/python/testing/recipe
     library/python/testing/yatest_common
diff --git a/ydb/tests/fq/generic/analytics/ydb/01_basic.sh b/ydb/tests/fq/generic/analytics/ydb/01_basic.sh
@@ -35,6 +35,9 @@ set -ex
       (56, 12, "2a02:1812:1713:4f00:517e:1d79:c88b:704", "Elena", 2),
       (18, 17, "ivalid ip", "newUser", 12);
     COMMIT;
+
+    CREATE TABLE dummy_table(name STRING, cnt Uint64, PRIMARY KEY(name));
+    COMMIT;
   '
 
 retVal=$?
diff --git a/ydb/tests/fq/generic/streaming/conftest.py b/ydb/tests/fq/generic/streaming/conftest.py
@@ -10,10 +10,14 @@
 
 from ydb.tests.fq.generic.utils.settings import Settings
 
+from typing import Final
+
+docker_compose_file_path: Final = "ydb/tests/fq/generic/streaming/docker-compose.yml"
+
 
 @pytest.fixture
 def settings() -> Settings:
-    return Settings.from_env(docker_compose_file_path='ydb/tests/fq/generic/streaming/docker-compose.yml')
+    return Settings.from_env(docker_compose_file_path=docker_compose_file_path)
 
 
 @pytest.fixture
diff --git a/ydb/tests/fq/generic/streaming/test_join.py b/ydb/tests/fq/generic/streaming/test_join.py
@@ -10,6 +10,8 @@
 
 from ydb.tests.tools.fq_runner.fq_client import FederatedQueryClient
 from ydb.tests.tools.datastreams_helpers.test_yds_base import TestYdsBase
+from ydb.library.yql.providers.generic.connector.tests.utils.scenario.ydb import OneTimeWaiter
+import conftest
 
 DEBUG = 0
 
@@ -395,6 +397,12 @@ def freeze(json):
 ]
 
 
+one_time_waiter = OneTimeWaiter(
+    docker_compose_file_path=conftest.docker_compose_file_path,
+    expected_tables=["simple_table", "join_table", "dummy_table"],
+)
+
+
 class TestJoinStreaming(TestYdsBase):
     @yq_v1
     @pytest.mark.parametrize(
@@ -412,6 +420,7 @@ def test_simple(self, kikimr, fq_client: FederatedQueryClient, yq_version):
             name=ydb_conn_name,
             database_id='local',
         )
+        one_time_waiter.wait()
 
         sql = R'''
             $input = SELECT * FROM myyds.`{input_topic}`;
@@ -487,6 +496,8 @@ def test_streamlookup(
             streamlookup=R'/*+ streamlookup() */' if streamlookup else '',
         )
 
+        one_time_waiter.wait()
+
         query_id = fq_client.create_query(
             f"streamlookup_{partitions_count}{streamlookup}{testcase}", sql, type=fq.QueryContent.QueryType.STREAMING
         ).result.query_id
diff --git a/ydb/tests/fq/generic/streaming/ya.make b/ydb/tests/fq/generic/streaming/ya.make
@@ -54,6 +54,7 @@ DEPENDS(
 )
 
 PEERDIR(
+    ydb/library/yql/providers/generic/connector/tests/utils/scenario
     ydb/tests/fq/generic/utils
     ydb/tests/tools/datastreams_helpers
     library/python/testing/recipe
diff --git a/ydb/tests/fq/generic/streaming/ydb/01_basic.sh b/ydb/tests/fq/generic/streaming/ydb/01_basic.sh
@@ -49,4 +49,17 @@ if [ $retVal -ne 0 ]; then
   exit $retVal
 fi
 
+/ydb -p ${PROFILE} yql -s '
+    CREATE TABLE dummy_table (name String, cnt Uint64, PRIMARY KEY(name));
+    COMMIT;
+    INSERT INTO dummy_table (name, cnt) SELECT "db", COUNT(*) FROM db;
+    COMMIT;
+  '
+
+retVal=$?
+if [ $retVal -ne 0 ]; then
+  echo $retVal
+  exit $retVal
+fi
+
 echo $(date +"%T.%6N") "SUCCESS"

Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,7 @@ PEERDIR(`
`20`	`20`	`ydb/library/yql/providers/generic/connector/tests/utils`
`21`	`21`	`ydb/library/yql/providers/generic/connector/tests/utils/clients`
`22`	`22`	`ydb/library/yql/providers/generic/connector/tests/utils/run`
	`23`	`+ ydb/library/yql/providers/generic/connector/tests/common_test_cases`
`23`	`24`	`)`
`24`	`25`
`25`	`26`	`END()`
Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,9 @@ set -ex`
`35`	`35`	`(56, 12, "2a02:1812:1713:4f00:517e:1d79:c88b:704", "Elena", 2),`
`36`	`36`	`(18, 17, "ivalid ip", "newUser", 12);`
`37`	`37`	`COMMIT;`
	`38`	`+`
	`39`	`+ CREATE TABLE dummy_table(name STRING, cnt Uint64, PRIMARY KEY(name));`
	`40`	`+ COMMIT;`
`38`	`41`	`'`
`39`	`42`
`40`	`43`	`retVal=$?`