fix tests in oss for new sdk version

4 files changed, 18 insertions, 6 deletions

diff --git a/ydb/public/sdk/python3/ydb/_errors.py b/ydb/public/sdk/python3/ydb/_errors.py
index ae3057b6d2f..8c6f072049f 100644
--- a/ydb/public/sdk/python3/ydb/_errors.py
+++ b/ydb/public/sdk/python3/ydb/_errors.py
@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Optional
 
-from ydb import issues
+from . import issues
 
 _errors_retriable_fast_backoff_types = [
     issues.Unavailable,
diff --git a/ydb/public/sdk/python3/ydb/scheme_test.py b/ydb/public/sdk/python3/ydb/scheme_test.py
index 79094882923..d63850bf930 100644
--- a/ydb/public/sdk/python3/ydb/scheme_test.py
+++ b/ydb/public/sdk/python3/ydb/scheme_test.py
@@ -1,9 +1,9 @@
-from ydb.scheme import (
+from .scheme import (
     SchemeEntryType,
     _wrap_scheme_entry,
     _wrap_list_directory_response,
 )
-from ydb._apis import ydb_scheme
+from ._apis import ydb_scheme
 
 
 def test_wrap_scheme_entry():
diff --git a/ydb/public/sdk/python3/ydb/table_test.py b/ydb/public/sdk/python3/ydb/table_test.py
index 2cb2a6a0d0d..8e93a698f4a 100644
--- a/ydb/public/sdk/python3/ydb/table_test.py
+++ b/ydb/public/sdk/python3/ydb/table_test.py
@@ -1,5 +1,5 @@
 from unittest import mock
-from ydb import (
+from . import (
     retry_operation_impl,
     YdbRetryOperationFinalResult,
     issues,
diff --git a/ydb/tests/oss/ydb_sdk_import/__init__.py b/ydb/tests/oss/ydb_sdk_import/__init__.py
index 31ff3cca1f2..e29b21910bf 100644
--- a/ydb/tests/oss/ydb_sdk_import/__init__.py
+++ b/ydb/tests/oss/ydb_sdk_import/__init__.py
@@ -1,9 +1,21 @@
 from ydb.tests.oss.canonical import is_oss
 import os
-
+import sys
 
 if is_oss:
-    if os.getenv('PYTHON2_YDB_IMPORT'):
+    python_2 = os.getenv('PYTHON2_YDB_IMPORT')
+    sdk_path = ""
+    if python_2:
+        sdk_path = "ydb.public.sdk.python2"
+    else:
+        sdk_path = "ydb.public.sdk.python3"
+
+    from ydb.public.api.grpc import *  # noqa
+    sys.modules[sdk_path + ".ydb._grpc.common"] = sys.modules["ydb.public.api.grpc"]
+    from ydb.public.api import protos  # noqa
+    sys.modules[sdk_path + ".ydb._grpc.common.protos"] = sys.modules["ydb.public.api.protos"]
+
+    if python_2:
         from ydb.public.sdk.python2 import ydb # noqa
     else:
         from ydb.public.sdk.python3 import ydb # noqa