add ut

Signed-off-by: luohaha <[email protected]>
StarRocks · Jul 18, 2024 · 3f5741f · 3f5741f
1 parent 75827f4
commit 3f5741f
Showing 1 changed file with 31 additions and 0 deletions.
diff --git a/be/test/storage/lake/delta_writer_test.cpp b/be/test/storage/lake/delta_writer_test.cpp
@@ -610,4 +610,35 @@ TEST_F(LakeDeltaWriterTest, test_memtable_full) {
     ASSERT_GT(txnlog->op_write().rowset().data_size(), 0);
 }
 
+TEST_F(LakeDeltaWriterTest, test_write_oom) {
+    // Prepare data for writing
+    static const int kChunkSize = 128;
+    auto chunk0 = generate_data(kChunkSize);
+    auto indexes = std::vector<uint32_t>(kChunkSize);
+    for (int i = 0; i < kChunkSize; i++) {
+        indexes[i] = i;
+    }
+
+    // Create and open DeltaWriter
+    auto txn_id = next_id();
+    auto tablet_id = _tablet_metadata->id();
+    int64_t old_limit = GlobalEnv::GetInstance()->load_mem_tracker()->limit();
+    GlobalEnv::GetInstance()->load_mem_tracker()->set_limit(1);
+    GlobalEnv::GetInstance()->load_mem_tracker()->consume(100);
+    ASSIGN_OR_ABORT(auto delta_writer, DeltaWriterBuilder()
+                                               .set_tablet_manager(_tablet_mgr.get())
+                                               .set_tablet_id(tablet_id)
+                                               .set_txn_id(txn_id)
+                                               .set_partition_id(_partition_id)
+                                               .set_mem_tracker(_mem_tracker.get())
+                                               .set_schema_id(_tablet_schema->id())
+                                               .set_immutable_tablet_size(1)
+                                               .build());
+    ASSERT_OK(delta_writer->open());
+    // Write and flush
+    ASSERT_ERROR(delta_writer->write(chunk0, indexes.data(), indexes.size()));
+    GlobalEnv::GetInstance()->load_mem_tracker()->release(100);
+    GlobalEnv::GetInstance()->load_mem_tracker()->set_limit(old_limit);
+}
+
 } // namespace starrocks::lake