Move kikimr/yq/tests/plans to ydb

4 files changed, 178 insertions, 0 deletions

diff --git a/ydb/tests/fq/plans/conftest.py b/ydb/tests/fq/plans/conftest.py
new file mode 100644
index 0000000000..786126d90d
--- /dev/null
+++ b/ydb/tests/fq/plans/conftest.py
@@ -0,0 +1,88 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import logging
+import pytest
+import requests
+import yatest.common
+from dataclasses import dataclass
+
+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 AddInflightExtension
+from ydb.tests.tools.fq_runner.kikimr_utils import AddDataInflightExtension
+from ydb.tests.tools.fq_runner.kikimr_utils import AddFormatSizeLimitExtension
+from ydb.tests.tools.fq_runner.kikimr_utils import DefaultConfigExtension
+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 StatsModeExtension
+from ydb.tests.tools.fq_runner.kikimr_utils import start_kikimr
+from library.recipes import common as recipes_common
+
+
+MOTO_SERVER_PATH = "contrib/python/moto/bin/moto_server"
+S3_PID_FILE = "s3.pid"
+
+
+@dataclass
+class S3:
+    s3_url: str
+
+
+@pytest.fixture(scope="module")
+def s3(request) -> S3:
+    port_manager = yatest.common.network.PortManager()
+    s3_port = port_manager.get_port()
+    s3_url = "http://localhost:{port}".format(port=s3_port)
+
+    command = [
+        yatest.common.binary_path(MOTO_SERVER_PATH),
+        "s3",
+        "--host", "::1",
+        "--port", str(s3_port)
+    ]
+
+    def is_s3_ready():
+        try:
+            response = requests.get(s3_url)
+            response.raise_for_status()
+            return True
+        except requests.RequestException as err:
+            logging.debug(err)
+            return False
+
+    recipes_common.start_daemon(
+        command=command,
+        environment=None,
+        is_alive_check=is_s3_ready,
+        pid_file_name=S3_PID_FILE
+    )
+
+    try:
+        yield S3(s3_url)
+    finally:
+        with open(S3_PID_FILE, 'r') as f:
+            pid = int(f.read())
+        recipes_common.stop_daemon(pid)
+
+
+@pytest.fixture
+def kikimr(request: pytest.FixtureRequest, s3: S3, yq_version: str, stats_mode: str):
+    kikimr_extensions = [AddInflightExtension(),
+                         AddDataInflightExtension(),
+                         AddFormatSizeLimitExtension(),
+                         DefaultConfigExtension(s3.s3_url),
+                         YQv2Extension(yq_version),
+                         ComputeExtension(),
+                         StatsModeExtension(stats_mode)]
+    with start_kikimr(request, kikimr_extensions) as kikimr:
+        yield kikimr
+
+
+@pytest.fixture
+def client(kikimr, request=None):
+    client = FederatedQueryClient(request.param["folder_id"]
+                                  if request is not None
+                                  else "my_folder",
+                                  streaming_over_kikimr=kikimr)
+    return client
diff --git a/ydb/tests/fq/plans/test_stats_mode.py b/ydb/tests/fq/plans/test_stats_mode.py
new file mode 100644
index 0000000000..46f56ade7b
--- /dev/null
+++ b/ydb/tests/fq/plans/test_stats_mode.py
@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import boto3
+import logging
+import pytest
+import ydb.public.api.protos.draft.fq_pb2 as fq
+from ydb.tests.tools.fq_runner.kikimr_utils import yq_all
+
+
+class TestStatsMode:
+
+    @yq_all
+    @pytest.mark.parametrize("client", [{"folder_id": "my_folder"}], indirect=True)
+    @pytest.mark.parametrize("stats_mode", ["STATS_MODE_NONE", "STATS_MODE_BASIC", "STATS_MODE_FULL", "STATS_MODE_PROFILE"])
+    def test_mode(self, kikimr, s3, client, stats_mode):
+        resource = boto3.resource(
+            "s3",
+            endpoint_url=s3.s3_url,
+            aws_access_key_id="key",
+            aws_secret_access_key="secret_key"
+        )
+
+        bucket = resource.Bucket("pbucket")
+        bucket.create(ACL='public-read')
+        kikimr.control_plane.wait_bootstrap(1)
+        client.create_storage_connection("pb", "pbucket")
+
+        sql = R'''
+            insert into pb.`path/` with (format=csv_with_names)
+            select * from AS_TABLE([<|foo:1, bar:"xxx"u|>,<|foo:2, bar:"yyy"u|>]);
+            '''
+        query_id = client.create_query("simple", sql, type=fq.QueryContent.QueryType.ANALYTICS).result.query_id
+        client.wait_query_status(query_id, fq.QueryMeta.COMPLETED)
+
+        sql = R'''
+            insert into pb.`path/` with (format=csv_with_names)
+            select * from AS_TABLE([<|foo:3, bar:"xxx"u|>,<|foo:4, bar:"yyy"u|>]);
+            '''
+        query_id = client.create_query("simple", sql, type=fq.QueryContent.QueryType.ANALYTICS).result.query_id
+        client.wait_query_status(query_id, fq.QueryMeta.COMPLETED)
+
+        sql = R'''
+            select bar, count(foo) as foo_count, sum(foo) as foo_sum
+            from pb.`path/` with (format=csv_with_names, schema(
+                foo Int NOT NULL,
+                bar String NOT NULL
+            ))
+            group by bar
+            '''
+
+        query_id = client.create_query("simple", sql, type=fq.QueryContent.QueryType.ANALYTICS).result.query_id
+        client.wait_query_status(query_id, fq.QueryMeta.COMPLETED)
+
+        data = client.get_result_data(query_id)
+        result_set = data.result.result_set
+        logging.debug(str(result_set))
+        assert len(result_set.columns) == 3
+        assert len(result_set.rows) == 2
+        # assert result_set.rows[0].items[0].uint64_value == 1024 * 10
+        # 1024 x 1024 x 10 = 10 MB of raw data + little overhead for header, eols etc
+        # assert sum(kikimr.control_plane.get_metering()) == 11
diff --git a/ydb/tests/fq/plans/ya.make b/ydb/tests/fq/plans/ya.make
new file mode 100644
index 0000000000..daf5fc7399
--- /dev/null
+++ b/ydb/tests/fq/plans/ya.make
@@ -0,0 +1,27 @@
+PY3TEST()
+
+INCLUDE(${ARCADIA_ROOT}/ydb/tests/tools/fq_runner/ydb_runner_with_datastreams.inc)
+
+PEERDIR(
+    contrib/python/boto3
+    library/python/testing/recipe
+    library/python/testing/yatest_common
+    library/recipes/common
+    ydb/tests/tools/fq_runner
+)
+
+DEPENDS(
+    contrib/python/moto/bin
+)
+
+TEST_SRCS(
+    test_stats_mode.py
+)
+
+PY_SRCS(
+    conftest.py
+)
+
+SIZE(MEDIUM)
+
+END()
diff --git a/ydb/tests/fq/ya.make b/ydb/tests/fq/ya.make
index 822c35f9e5..77586532cc 100644
--- a/ydb/tests/fq/ya.make
+++ b/ydb/tests/fq/ya.make
@@ -1,3 +1,4 @@
 RECURSE_FOR_TESTS(
+    plans
     s3
 )