Add test

3 files changed, 24 insertions, 2 deletions

diff --git a/ydb/core/base/appdata.h b/ydb/core/base/appdata.h
index ae764fdf9bc..d70f2955b1c 100644
--- a/ydb/core/base/appdata.h
+++ b/ydb/core/base/appdata.h
@@ -165,7 +165,9 @@ struct TAppData {
 
     // Used to disable object deletion in schemeshard for cleanup tests
     bool DisableSchemeShardCleanupOnDropForTest = false;
-    
+
+    TMaybe<ui32> ZstdBlockSizeForTest;
+
     // Used to disable checking nodes with sys tablets only in cms
     bool  DisableCheckingSysNodesCms = false;
 
diff --git a/ydb/core/tx/datashard/import_s3.cpp b/ydb/core/tx/datashard/import_s3.cpp
index 4599eddd246..1fa6100667b 100644
--- a/ydb/core/tx/datashard/import_s3.cpp
+++ b/ydb/core/tx/datashard/import_s3.cpp
@@ -189,7 +189,7 @@ class TS3Downloader: public TActorBootstrapped<TS3Downloader> {
                         return ERROR;
                     }
 
-                    Buffer.Reserve(Buffer.Capacity() + ZSTD_BLOCKSIZE_MAX);
+                    Buffer.Reserve(Buffer.Capacity() + AppData()->ZstdBlockSizeForTest.GetOrElse(ZSTD_BLOCKSIZE_MAX));
                 }
 
                 PendingInputBytes += input.pos; // inc after decompress
diff --git a/ydb/core/tx/schemeshard/ut_restore.cpp b/ydb/core/tx/schemeshard/ut_restore.cpp
index 422f4cb6a8c..c1e05806b3f 100644
--- a/ydb/core/tx/schemeshard/ut_restore.cpp
+++ b/ydb/core/tx/schemeshard/ut_restore.cpp
@@ -446,6 +446,26 @@ Y_UNIT_TEST_SUITE(TRestoreTests) {
         ShouldSucceedOnMultipleFrames(1);
     }
 
+    Y_UNIT_TEST(ShouldSucceedOnSmallBuffer) {
+        TTestBasicRuntime runtime;
+        TTestEnv env(runtime);
+        runtime.GetAppData().ZstdBlockSizeForTest = 16;
+        runtime.GetAppData().DataShardConfig.SetRestoreReadBufferSizeLimit(16);
+
+        const auto data = GenerateZstdTestData("a", 2);
+        const ui32 batchSize = 1;
+
+        Restore(runtime, env, R"(
+            Name: "Table"
+            Columns { Name: "key" Type: "Utf8" }
+            Columns { Name: "value" Type: "Utf8" }
+            KeyColumnNames: ["key"]
+        )", {data}, batchSize);
+
+        auto content = ReadTable(runtime, TTestTxConfig::FakeHiveTablets);
+        NKqp::CompareYson(data.YsonStr, content);
+    }
+
     Y_UNIT_TEST_WITH_COMPRESSION(ShouldExpandBuffer) {
         TTestBasicRuntime runtime;