Segment test framework (#5150)

close #5151

dbms/src/Storages/DeltaMerge/tests/DMTestEnv.h

@@ -273,7 +273,8 @@ class DMTestEnv
                                          DataTypePtr pk_type = EXTRA_HANDLE_COLUMN_INT_TYPE,
                                          bool is_common_handle = false,
                                          size_t rowkey_column_size = 1,
-                                         bool with_internal_columns = true)
+                                         bool with_internal_columns = true,
+                                         bool is_deleted = false)
     {
         Block block;
         const size_t num_rows = (end - beg);
@@ -324,7 +325,7 @@ class DMTestEnv
                 VERSION_COLUMN_ID));
             // tag_col
             block.insert(DB::tests::createColumn<UInt8>(
-                std::vector<UInt64>(num_rows, 0),
+                std::vector<UInt64>(num_rows, is_deleted),
                 TAG_COLUMN_NAME,
                 TAG_COLUMN_ID));
         }

dbms/src/Storages/DeltaMerge/tests/gtest_segment.cpp

@@ -0,0 +1,86 @@
+// Copyright 2022 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include <Common/CurrentMetrics.h>
+#include <DataStreams/OneBlockInputStream.h>
+#include <Storages/DeltaMerge/DeltaMergeStore.h>
+#include <Storages/DeltaMerge/tests/gtest_segment_test_basic.h>
+#include <TestUtils/TiFlashTestBasic.h>
+
+
+namespace DB
+{
+namespace DM
+{
+namespace tests
+{
+class SegmentOperationTest : public SegmentTestBasic
+{
+protected:
+    static void SetUpTestCase() {}
+};
+
+TEST_F(SegmentOperationTest, Issue4956)
+try
+{
+    SegmentTestOptions options;
+    reloadWithOptions(options);
+
+    // flush data, make the segment can be split.
+    writeSegment(DELTA_MERGE_FIRST_SEGMENT_ID);
+    flushSegmentCache(DELTA_MERGE_FIRST_SEGMENT_ID);
+    // write data to cache, reproduce the https://github.com/pingcap/tiflash/issues/4956
+    writeSegment(DELTA_MERGE_FIRST_SEGMENT_ID);
+    deleteRangeSegment(DELTA_MERGE_FIRST_SEGMENT_ID);
+    auto segment_id = splitSegment(DELTA_MERGE_FIRST_SEGMENT_ID);
+    ASSERT_TRUE(segment_id.has_value());
+
+    mergeSegment(DELTA_MERGE_FIRST_SEGMENT_ID, *segment_id);
+}
+CATCH
+
+TEST_F(SegmentOperationTest, TestSegment)
+try
+{
+    SegmentTestOptions options;
+    reloadWithOptions(options);
+    writeSegment(DELTA_MERGE_FIRST_SEGMENT_ID);
+    flushSegmentCache(DELTA_MERGE_FIRST_SEGMENT_ID);
+    mergeSegmentDelta(DELTA_MERGE_FIRST_SEGMENT_ID);
+    auto segment_id = splitSegment(DELTA_MERGE_FIRST_SEGMENT_ID);
+    ASSERT_TRUE(segment_id.has_value());
+
+    size_t origin_rows = getSegmentRowNum(DELTA_MERGE_FIRST_SEGMENT_ID);
+
+    writeSegment(*segment_id);
+    flushSegmentCache(*segment_id);
+    deleteRangeSegment(*segment_id);
+    writeSegmentWithDeletedPack(*segment_id);
+    mergeSegment(DELTA_MERGE_FIRST_SEGMENT_ID, *segment_id);
+
+    EXPECT_EQ(getSegmentRowNum(DELTA_MERGE_FIRST_SEGMENT_ID), origin_rows);
+}
+CATCH
+
+TEST_F(SegmentOperationTest, TestSegmentRandom)
+try
+{
+    SegmentTestOptions options;
+    options.is_common_handle = true;
+    reloadWithOptions(options);
+    randomSegmentTest(100);
+}
+CATCH
+} // namespace tests
+} // namespace DM
+} // namespace DB