Skip to content

Commit

Permalink
test
Browse files Browse the repository at this point in the history
  • Loading branch information
vlad-gogov committed Nov 29, 2024
1 parent 0392fbf commit 98db446
Show file tree
Hide file tree
Showing 3 changed files with 82 additions and 0 deletions.
20 changes: 20 additions & 0 deletions ydb/tests/olap/scenario/helpers/scenario_tests_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,26 @@ def execute_scan_query(
allure.attach(json.dumps(rows), 'result', allure.attachment_type.JSON)
return ret

@allure.step('Execute query')
def execute_query(
self, yql: str, expected_status: ydb.StatusCode | Set[ydb.StatusCode] = ydb.StatusCode.SUCCESS
):
"""Run a query on the tested database.
Args:
yql: Query text.
expected_status: Expected status or set of database response statuses. If the response status is not in the expected set, an exception is thrown.
Example:
tablename = 'testTable'
sth = ScenarioTestHelper(ctx)
sth.execute_query(f'INSERT INTO `{sth.get_full_path("tablename") }` (key, c) values(1, 100)')
"""

allure.attach(yql, 'request', allure.attachment_type.TEXT)
with ydb.QuerySessionPool(YdbCluster.get_ydb_driver()) as pool:
it = self._run_with_expected_status(lambda: pool.execute_with_retries(yql), expected_status)

def drop_if_exist(self, names: List[str], operation) -> None:
"""Erase entities in the tested database, if it exists.
Expand Down
61 changes: 61 additions & 0 deletions ydb/tests/olap/scenario/test_insert.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
from conftest import BaseTestSet
from ydb.tests.olap.scenario.helpers import (
ScenarioTestHelper,
TestContext,
CreateTable,
DropTable,
)

from ydb import PrimitiveType
from typing import List, Dict, Any
import threading

class TestInsert(BaseTestSet):
schema_cnt = (
ScenarioTestHelper.Schema()
.with_column(name='key', type=PrimitiveType.Int32, not_null=True)
.with_column(name='c', type=PrimitiveType.Int64)
.with_key_columns('key')
)

schema_log = (
ScenarioTestHelper.Schema()
.with_column(name='key', type=PrimitiveType.Int32, not_null=True)
.with_key_columns('key')
)

def _loop_upsert(self, ctx: TestContext, data: list):
sth = ScenarioTestHelper(ctx)
for batch in data:
sth.bulk_upsert_data("log", self.schema_log, batch)

def _loop_insert(self, ctx: TestContext):
sth = ScenarioTestHelper(ctx)
for i in range(100):
sth.execute_query(f"$cnt = SELECT CAST(COUNT(*) AS INT64) from `{sth.get_full_path("log")}`; INSERT INTO `{sth.get_full_path("cnt") }` (key, c) values({i}, $cnt)")

def scenario_read_data_during_bulk_upsert(self, ctx: TestContext):
sth = ScenarioTestHelper(ctx)
cnt_table_name: str = "cnt"
log_table_name: str = "log"
sth.execute_scheme_query(CreateTable(cnt_table_name).with_schema(self.schema_cnt))
sth.execute_scheme_query(CreateTable(log_table_name).with_schema(self.schema_log))

data: list = []
for i in range(100):
batch: List[Dict[str, Any]] = []
for j in range(i, 1000):
batch.append({'key' : j})
data.append(batch)

thread1 = threading.Thread(target=self._loop_upsert, args=[ctx, data])
thread2 = threading.Thread(target=self._loop_insert, args=[ctx])

thread1.start()
thread2.start()

thread2.join()
thread1.join()

sth.execute_scheme_query(DropTable(cnt_table_name))
sth.execute_scheme_query(DropTable(log_table_name))
1 change: 1 addition & 0 deletions ydb/tests/olap/scenario/ya.make
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ PY3TEST()
test_simple.py
test_scheme_load.py
test_alter_tiering.py
test_insert.py
)

PEERDIR(
Expand Down

0 comments on commit 98db446

Please sign in to comment.