Switch tests to the new tablet service (#8454)

6 files changed, 91 insertions, 65 deletions

diff --git a/ydb/tests/acceptance/cluster.yaml b/ydb/tests/acceptance/cluster.yaml
index 5c3683dc5b..547ea3e1e1 100644
--- a/ydb/tests/acceptance/cluster.yaml
+++ b/ydb/tests/acceptance/cluster.yaml
@@ -117,6 +117,7 @@ grpc:
   - cms
   - discovery
   - legacy
+  - tablet_service
   - experimental
   - pq
   - clickhouse_internal
diff --git a/ydb/tests/functional/hive/test_kill_tablets.py b/ydb/tests/functional/hive/test_kill_tablets.py
index 08dbd5c0b1..b0fb1f2557 100644
--- a/ydb/tests/functional/hive/test_kill_tablets.py
+++ b/ydb/tests/functional/hive/test_kill_tablets.py
@@ -4,7 +4,7 @@ from hamcrest import assert_that, greater_than, has_length
 from ydb.tests.library.common.delayed import wait_tablets_state_by_id
 from ydb.tests.library.common.types import TabletTypes, TabletStates
 from ydb.tests.library.harness.kikimr_cluster import kikimr_cluster_factory
-from ydb.tests.library.matchers.response import is_ok_response, is_valid_response_with_field
+from ydb.tests.library.matchers.response import is_valid_response_with_field
 from ydb.tests.library.kv.helpers import create_tablets_and_wait_for_start
 
 TIMEOUT_SECONDS = 180
@@ -51,11 +51,7 @@ class TestKillTablets(object):
         hive_id = hive_state_response.TabletStateInfo[0].TabletId
 
         # Act
-        response = self.cluster.client.tablet_kill(hive_id)
-        assert_that(
-            response,
-            is_ok_response()
-        )
+        self.cluster.client.tablet_kill(hive_id, assert_success=True)
 
         wait_tablets_state_by_id(
             self.cluster.client,
diff --git a/ydb/tests/library/common/local_db_scheme.py b/ydb/tests/library/common/local_db_scheme.py
index 2350333f6c..42a545b892 100644
--- a/ydb/tests/library/common/local_db_scheme.py
+++ b/ydb/tests/library/common/local_db_scheme.py
@@ -1,15 +1,17 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-import ydb.core.protos.msgbus_pb2 as msgbus
-from ydb.tests.library.common.types import DeltaTypes, PType
+import json
+from ydb.public.api.protos.ydb_status_codes_pb2 import StatusIds
+from ydb.public.api.protos.draft.ydb_tablet_pb2 import ChangeTabletSchemaRequest
+from ydb.tests.library.common.types import PType
 
 
 def column_to_dict(delta):
     return {
-        "ColumnId": delta.ColumnId,
-        "ColumnName": delta.ColumnName,
+        "ColumnId": delta['column_id'],
+        "ColumnName": delta['column_name'],
         "ColumnType": str(
-            PType.from_int(delta.ColumnType)
+            PType.from_int(delta['column_type'])
         )
     }
 
@@ -35,8 +37,8 @@ class TableScheme(object):
         return self
 
     def add_column_to_family(self, delta):
-        cid = delta.ColumnId
-        fid = delta.FamilyId
+        cid = delta['column_id']
+        fid = delta['family_id']
 
         if fid not in self.data['ColumnFamilies']:
             raise RuntimeError("Unknown family")
@@ -45,8 +47,8 @@ class TableScheme(object):
         return self
 
     def add_family(self, delta):
-        fid = delta.FamilyId
-        rid = delta.RoomId
+        fid = delta['family_id']
+        rid = delta['room_id']
 
         if rid not in self.data['Rooms']:
             raise RuntimeError("Unknown room")
@@ -67,25 +69,25 @@ class TableScheme(object):
         return self
 
     def set_family(self, delta):
-        fid = delta.FamilyId
+        fid = delta['family_id']
 
         if fid not in self.data['ColumnFamilies']:
             raise RuntimeError("Unknown family")
 
-        if hasattr(delta, 'Codec') and delta.Codec is not None:
-            self.data['ColumnFamilies'][fid]['Codec'] = delta.Codec
+        if 'codec' in delta and delta['codec'] is not None:
+            self.data['ColumnFamilies'][fid]['Codec'] = delta['codec']
 
-        if hasattr(delta, 'InMemory') and delta.InMemory is not None:
-            self.data['ColumnFamilies'][fid]['InMemory'] = delta.InMemory
+        if 'in_memory' in delta and delta['in_memory'] is not None:
+            self.data['ColumnFamilies'][fid]['InMemory'] = delta['in_memory']
 
-        if hasattr(delta, 'Cache') and delta.Cache is not None:
-            self.data['ColumnFamilies'][fid]['Cache'] = delta.Cache
+        if 'cache' in delta and delta['cache'] is not None:
+            self.data['ColumnFamilies'][fid]['Cache'] = delta['cache']
 
-        if hasattr(delta, 'Small') and delta.Small is not None:
-            self.data['ColumnFamilies'][fid]['Small'] = delta.Small
+        if 'small' in delta and delta['small'] is not None:
+            self.data['ColumnFamilies'][fid]['Small'] = delta['small']
 
-        if hasattr(delta, 'Large') and delta.Large is not None:
-            self.data['ColumnFamilies'][fid]['Large'] = delta.Large
+        if 'large' in delta and delta['large'] is not None:
+            self.data['ColumnFamilies'][fid]['Large'] = delta['large']
 
         return self
 
@@ -102,7 +104,7 @@ class TableScheme(object):
         return self.data['Rooms']
 
     def add_column_to_key(self, delta):
-        cid = delta.ColumnId
+        cid = delta['column_id']
         self.data['TableKey'].append(cid)
         return self
 
@@ -110,7 +112,7 @@ class TableScheme(object):
         raise RuntimeError("Unknown delta type")
 
     def set_room(self, delta):
-        self.data['Rooms'][delta.RoomId] = {"Main": delta.Main, "Outer": delta.Outer, "Blobs": delta.Blobs}
+        self.data['Rooms'][delta['room_id']] = {"Main": delta['main'], "Outer": delta['outer'], "Blobs": delta['blobs']}
         return self
 
     def set_compaction_policy(self, delta):
@@ -123,23 +125,27 @@ class TableScheme(object):
 
     def add_delta(self, delta):
         mapping = {
-            DeltaTypes.AddColumn: self.add_column,
-            DeltaTypes.DropColumn: self.drop_column,
-            DeltaTypes.AddColumnToKey: self.add_column_to_key,
-            DeltaTypes.AddColumnToFamily: self.add_column_to_family,
-            DeltaTypes.AddFamily: self.add_family,
-            DeltaTypes.SetRoom: self.set_room,
-            DeltaTypes.SetCompactionPolicy: self.set_compaction_policy,
-            DeltaTypes.SetFamily: self.set_family,
-            DeltaTypes.SetTable: self.set_table,
+            'AddColumn': self.add_column,
+            'DropColumn': self.drop_column,
+            'AddColumnToKey': self.add_column_to_key,
+            'AddColumnToFamily': self.add_column_to_family,
+            'AddFamily': self.add_family,
+            'SetRoom': self.set_room,
+            'SetCompactionPolicy': self.set_compaction_policy,
+            'SetFamily': self.set_family,
+            'SetTable': self.set_table,
         }
-        op = mapping.get(delta.DeltaType, self.unknown_delta)
+        op = mapping.get(delta['delta_type'], self.unknown_delta)
         return op(delta)
 
 
 def get_deltas(client, tablet_id):
-    resp = client.send_request(msgbus.TLocalSchemeTx(TabletID=tablet_id, Timeout=60 * 1000), method='LocalSchemeTx')
-    return resp.LocalDbScheme.Delta
+    resp = client.tablet_service.ChangeTabletSchema(ChangeTabletSchemaRequest(tablet_id=tablet_id))
+    assert resp.status == StatusIds.SUCCESS
+    if resp.status != StatusIds.SUCCESS:
+        raise RuntimeError('ERROR: {status} {issues}'.format(status=resp.status, issues=resp.issues))
+    schema = json.loads(resp.schema)
+    return schema['delta']
 
 
 def get_scheme(client, tablet_id):
@@ -148,15 +154,15 @@ def get_scheme(client, tablet_id):
     pos = {}
 
     for delta in deltas:
-        if not delta.HasField('TableId'):
+        if 'table_id' not in delta:
             continue
 
-        if delta.DeltaType == DeltaTypes.AddTable:
-            table = TableScheme(delta.TableId, delta.TableName)
-            pos[delta.TableId] = len(scheme)
+        if delta['delta_type'] == 'AddTable':
+            table = TableScheme(delta['table_id'], delta['table_name'])
+            pos[delta['table_id']] = len(scheme)
             scheme.append(table)
             continue
 
-        scheme[pos[delta.TableId]] = scheme[pos[delta.TableId]].add_delta(delta)
+        scheme[pos[delta['table_id']]] = scheme[pos[delta['table_id']]].add_delta(delta)
 
     return [element for element in scheme]
diff --git a/ydb/tests/library/common/types.py b/ydb/tests/library/common/types.py
index 5ec564d65f..05d023c35e 100644
--- a/ydb/tests/library/common/types.py
+++ b/ydb/tests/library/common/types.py
@@ -8,23 +8,6 @@ from ydb.tests.library.common.generators import int_between, one_of, float_in, s
 
 
 @unique
-class DeltaTypes(IntEnum):
-    AddTable = 1,
-    DropTable = 2,
-    AddColumn = 3,
-    DropColumn = 4,
-    AddColumnToKey = 5,
-    AddColumnToFamily = 6,
-    AddFamily = 7,
-    UpdateExecutorInfo = 8,
-    SetCompactionPolicy = 9,
-    SetRoom = 10,
-    SetFamily = 11,
-    SetRedo = 12,
-    SetTable = 13
-
-
-@unique
 class PDiskCategory(IntEnum):
     ROT = 0
     SSD = 1
diff --git a/ydb/tests/library/harness/kikimr_client.py b/ydb/tests/library/harness/kikimr_client.py
index 9c123c75de..4012437942 100644
--- a/ydb/tests/library/harness/kikimr_client.py
+++ b/ydb/tests/library/harness/kikimr_client.py
@@ -8,6 +8,7 @@ from ydb.tests.library.common.protobuf_ss import TSchemeOperationStatus
 
 import grpc
 import six
+import functools
 
 from google.protobuf.text_format import Parse
 from ydb.core.protos import blobstorage_config_pb2
@@ -17,6 +18,8 @@ import ydb.core.protos.flat_scheme_op_pb2 as flat_scheme_op_pb2
 import ydb.core.protos.grpc_pb2_grpc as grpc_server
 from ydb.core.protos import flat_scheme_op_pb2 as flat_scheme_op
 from ydb.public.api.protos.ydb_status_codes_pb2 import StatusIds
+from ydb.public.api.grpc.draft import ydb_tablet_v1_pb2_grpc as grpc_tablet_service
+from ydb.public.api.protos.draft.ydb_tablet_pb2 import RestartTabletRequest
 from collections import namedtuple
 
 
@@ -52,6 +55,36 @@ def to_bytes(v):
     return v
 
 
+class StubWithRetries(object):
+    __slots__ = ('_stub', '_retry_count', '_retry_min_sleep', '_retry_max_sleep', '__dict__')
+
+    def __init__(self, stub, retry_count=4, retry_min_sleep=0.1, retry_max_sleep=5):
+        self._stub = stub
+        self._retry_count = retry_count
+        self._retry_min_sleep = retry_min_sleep
+        self._retry_max_sleep = retry_max_sleep
+
+    def __getattr__(self, method):
+        target = getattr(self._stub, method)
+
+        @functools.wraps(target)
+        def wrapper(*args, **kwargs):
+            retries = self._retry_count
+            next_sleep = self._retry_min_sleep
+            while True:
+                try:
+                    return target(*args, **kwargs)
+                except (RuntimeError, grpc.RpcError):
+                    retries -= 1
+                    if retries <= 0:
+                        raise
+                    time.sleep(next_sleep)
+                    next_sleep = min(next_sleep * 2, self._retry_max_sleep)
+
+        setattr(self, method, wrapper)
+        return wrapper
+
+
 class KiKiMRMessageBusClient(object):
     def __init__(self, server, port, cluster=None, retry_count=1):
         self.server = server
@@ -66,6 +99,7 @@ class KiKiMRMessageBusClient(object):
         ]
         self._channel = grpc.insecure_channel("%s:%s" % (self.server, self.port), options=self._options)
         self._stub = grpc_server.TGRpcServerStub(self._channel)
+        self.tablet_service = StubWithRetries(grpc_tablet_service.TabletServiceStub(self._channel))
 
     def describe(self, path, token):
         request = msgbus.TSchemeDescribe()
@@ -350,9 +384,14 @@ class KiKiMRMessageBusClient(object):
             request.DeadlineInstantMs = deadline_ms
         return self.invoke(request, 'KeyValue')
 
-    def tablet_kill(self, tablet_id):
-        request = msgbus.TTabletKillRequest(TabletID=tablet_id)
-        return self.invoke(request, 'TabletKillRequest')
+    def tablet_kill(self, tablet_id, assert_success=False):
+        request = RestartTabletRequest(tablet_id=tablet_id)
+        response = self.tablet_service.RestartTablet(request)
+        if assert_success:
+            assert response.status == StatusIds.SUCCESS
+            if response.status != StatusIds.SUCCESS:
+                raise RuntimeError('ERROR: {status} {issues}'.format(status=response.status, issues=response.issues))
+        return response
 
     def tablet_state(self, tablet_type=None, tablet_ids=()):
         request = msgbus.TTabletStateRequest()
diff --git a/ydb/tests/library/harness/resources/default_yaml.yml b/ydb/tests/library/harness/resources/default_yaml.yml
index 7090fad7de..22e575930b 100644
--- a/ydb/tests/library/harness/resources/default_yaml.yml
+++ b/ydb/tests/library/harness/resources/default_yaml.yml
@@ -134,6 +134,7 @@ grpc_config:
   host: "{ydb_grpc_host}"
   services:
   - "legacy"
+  - "tablet_service"
   - "yql"
   - "discovery"
   - "cms"