Merge 551be8e into a9e6ce0

Author: kardymonds

ydb/tests/fq/kikimr/conftest.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import pytest
+
+from ydb.tests.tools.fq_runner.fq_client import FederatedQueryClient
+from ydb.tests.tools.fq_runner.custom_hooks import *  # noqa: F401,F403 Adding custom hooks for YQv2 support
+from ydb.tests.tools.fq_runner.kikimr_utils import ExtensionPoint
+from ydb.tests.tools.fq_runner.kikimr_utils import YQv2Extension
+from ydb.tests.tools.fq_runner.kikimr_utils import ComputeExtension
+from ydb.tests.tools.fq_runner.kikimr_utils import DefaultConfigExtension
+from ydb.tests.tools.fq_runner.kikimr_utils import StatsModeExtension
+from ydb.tests.tools.fq_runner.kikimr_utils import start_kikimr
+
+
+@pytest.fixture
+def stats_mode():
+    return ''
+
+
+@pytest.fixture
+def kikimr(request: pytest.FixtureRequest, yq_version: str, stats_mode: str):
+    kikimr_extensions = [DefaultConfigExtension(""),
+                         YQv2Extension(yq_version),
+                         ComputeExtension(),
+                         StatsModeExtension(stats_mode)]
+    with start_kikimr(request, kikimr_extensions) as kikimr:
+        yield kikimr
+
+
+class ManyRetriesConfigExtension(ExtensionPoint):
+    def __init__(self):
+        super().__init__()
+
+    def is_applicable(self, request):
+        return True
+
+    def apply_to_kikimr(self, request, kikimr):
+        kikimr.compute_plane.fq_config['control_plane_storage']['retry_policy_mapping'] = [
+            {
+                'status_code': [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
+                'policy': {
+                    'retry_count': 10000
+                }
+            }
+        ]
+
+
+@pytest.fixture
+def kikimr_many_retries(request: pytest.FixtureRequest, yq_version: str):
+    kikimr_extensions = [DefaultConfigExtension(""),
+                         ManyRetriesConfigExtension(),
+                         YQv2Extension(yq_version),
+                         ComputeExtension()]
+    with start_kikimr(request, kikimr_extensions) as kikimr:
+        yield kikimr
+
+
+def create_client(kikimr, request):
+    return FederatedQueryClient(request.param["folder_id"] if request is not None else "my_folder",
+                                streaming_over_kikimr=kikimr)
+
+
+@pytest.fixture
+def client(kikimr, request=None):
+    return create_client(kikimr, request)
+
+
+@pytest.fixture
+def client_many_retries(kikimr_many_retries, request=None):
+    return create_client(kikimr_many_retries, request)

ydb/tests/fq/kikimr/test_base.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from ydb.tests.tools.fq_runner.kikimr_runner import StreamingOverKikimr
+from ydb.tests.tools.fq_runner.kikimr_runner import StreamingOverKikimrConfig
+from ydb.tests.tools.datastreams_helpers.test_yds_base import TestYdsBase
+
+
+class TestBaseWithAbortingConfigParams(TestYdsBase):
+
+    @classmethod
+    def setup_class(cls):
+        kikimr_conf = StreamingOverKikimrConfig(cloud_mode=True)
+        cls.streaming_over_kikimr = StreamingOverKikimr(kikimr_conf)
+        cls.streaming_over_kikimr.control_plane.fq_config['control_plane_storage']['task_lease_ttl'] = "2s"
+        cls.streaming_over_kikimr.control_plane.fq_config['control_plane_storage']['task_lease_retry_policy'] = {}
+        cls.streaming_over_kikimr.control_plane.fq_config['control_plane_storage']['task_lease_retry_policy']['retry_count'] = 1
+        cls.streaming_over_kikimr.compute_plane.fq_config['pinger']['ping_period'] = "1s"
+        cls.streaming_over_kikimr.start_mvp_mock_server()
+        cls.streaming_over_kikimr.start()
+
+    @classmethod
+    def teardown_class(cls):
+        if hasattr(cls, "streaming_over_kikimr"):
+            cls.streaming_over_kikimr.stop_mvp_mock_server()
+            cls.streaming_over_kikimr.stop()