refine test utils

risingwavelabs · Jan 16, 2025 · 50f8a9d · 50f8a9d
1 parent df4360d
commit 50f8a9d
Show file tree

Hide file tree

Showing 12 changed files with 17 additions and 38 deletions.
diff --git a/src/batch/executors/src/executor/insert.rs b/src/batch/executors/src/executor/insert.rs
@@ -274,7 +274,7 @@ mod tests {
     use risingwave_common::transaction::transaction_message::TxnMsg;
     use risingwave_common::types::StructType;
     use risingwave_dml::dml_manager::DmlManager;
-    use risingwave_storage::hummock::test_utils::StateStoreReadTestExt;
+    use risingwave_storage::hummock::test_utils::*;
     use risingwave_storage::hummock::CachePolicy;
     use risingwave_storage::memory::MemoryStateStore;
     use risingwave_storage::store::ReadOptions;

diff --git a/src/ctl/src/cmd_impl/hummock/list_kv.rs b/src/ctl/src/cmd_impl/hummock/list_kv.rs
@@ -45,7 +45,7 @@ pub async fn list_kv(
     let range = (Unbounded, Unbounded);
     let read_snapshot = hummock
         .new_read_snapshot(
-            HummockReadEpoch::NoWait(epoch),
+            HummockReadEpoch::Committed(epoch),
             NewReadSnapshotOptions {
                 table_id: TableId { table_id },
             },

diff --git a/src/storage/hummock_test/benches/bench_hummock_iter.rs b/src/storage/hummock_test/benches/bench_hummock_iter.rs
@@ -28,7 +28,7 @@ use risingwave_hummock_test::test_utils::TestIngestBatch;
 use risingwave_meta::hummock::test_utils::setup_compute_env;
 use risingwave_meta::hummock::MockHummockMetaClient;
 use risingwave_storage::hummock::iterator::test_utils::mock_sstable_store;
-use risingwave_storage::hummock::test_utils::{default_opts_for_test, StateStoreReadTestExt};
+use risingwave_storage::hummock::test_utils::*;
 use risingwave_storage::hummock::{CachePolicy, HummockStorage};
 use risingwave_storage::storage_value::StorageValue;
 use risingwave_storage::store::*;

diff --git a/src/storage/hummock_test/src/bin/replay/replay_impl.rs b/src/storage/hummock_test/src/bin/replay/replay_impl.rs
@@ -32,10 +32,9 @@ use risingwave_pb::common::WorkerNode;
 use risingwave_pb::meta::subscribe_response::{Info, Operation as RespOperation};
 use risingwave_pb::meta::{SubscribeResponse, SubscribeType};
 use risingwave_storage::hummock::store::LocalHummockStorage;
+use risingwave_storage::hummock::test_utils::*;
 use risingwave_storage::hummock::HummockStorage;
-use risingwave_storage::store::{
-    to_owned_item, LocalStateStore, NewReadSnapshotOptions, StateStoreIterExt, StateStoreRead,
-};
+use risingwave_storage::store::{to_owned_item, LocalStateStore, StateStoreIterExt};
 use risingwave_storage::{StateStore, StateStoreIter, StateStoreReadIter};
 use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver};
 
@@ -112,18 +111,11 @@ impl ReplayRead for GlobalReplayImpl {
             key_range.1.map(TracedBytes::into).map(TableKey),
         );
 
-        let snapshot = self
+        let iter = self
             .store
-            .new_read_snapshot(
-                HummockReadEpoch::NoWait(epoch),
-                NewReadSnapshotOptions {
-                    table_id: read_options.table_id.table_id.into(),
-                },
-            )
+            .iter(key_range, epoch, read_options.into())
             .await
             .unwrap();
-
-        let iter = snapshot.iter(key_range, read_options.into()).await.unwrap();
         let stream = GlobalReplayIter::new(iter).into_stream().boxed();
         Ok(stream)
     }
@@ -134,18 +126,9 @@ impl ReplayRead for GlobalReplayImpl {
         epoch: u64,
         read_options: TracedReadOptions,
     ) -> Result<Option<TracedBytes>> {
-        let snapshot = self
+        Ok(self
             .store
-            .new_read_snapshot(
-                HummockReadEpoch::NoWait(epoch),
-                NewReadSnapshotOptions {
-                    table_id: read_options.table_id.table_id.into(),
-                },
-            )
-            .await
-            .unwrap();
-        Ok(snapshot
-            .get(TableKey(key.into()), read_options.into())
+            .get(TableKey(key.into()), epoch, read_options.into())
             .await
             .unwrap()
             .map(TracedBytes::from))

diff --git a/src/storage/hummock_test/src/compactor_tests.rs b/src/storage/hummock_test/src/compactor_tests.rs
@@ -69,7 +69,7 @@ pub(crate) mod tests {
         ConcatIterator, SkipWatermarkIterator, UserIterator,
     };
     use risingwave_storage::hummock::sstable_store::SstableStoreRef;
-    use risingwave_storage::hummock::test_utils::{gen_test_sstable_info, StateStoreReadTestExt};
+    use risingwave_storage::hummock::test_utils::*;
     use risingwave_storage::hummock::value::HummockValue;
     use risingwave_storage::hummock::{
         BlockedXor16FilterBuilder, CachePolicy, CompressionAlgorithm, FilterBuilder,

diff --git a/src/storage/hummock_test/src/failpoint_tests.rs b/src/storage/hummock_test/src/failpoint_tests.rs
@@ -26,9 +26,7 @@ use risingwave_meta::hummock::test_utils::setup_compute_env;
 use risingwave_meta::hummock::MockHummockMetaClient;
 use risingwave_rpc_client::HummockMetaClient;
 use risingwave_storage::hummock::iterator::test_utils::mock_sstable_store;
-use risingwave_storage::hummock::test_utils::{
-    count_stream, default_opts_for_test, StateStoreReadTestExt,
-};
+use risingwave_storage::hummock::test_utils::*;
 use risingwave_storage::hummock::{CachePolicy, HummockStorage};
 use risingwave_storage::storage_value::StorageValue;
 use risingwave_storage::store::{

diff --git a/src/storage/hummock_test/src/hummock_read_version_tests.rs b/src/storage/hummock_test/src/hummock_read_version_tests.rs
@@ -36,7 +36,7 @@ use risingwave_storage::hummock::shared_buffer::shared_buffer_batch::SharedBuffe
 use risingwave_storage::hummock::store::version::{
     read_filter_for_version, HummockReadVersion, StagingData, StagingSstableInfo, VersionUpdate,
 };
-use risingwave_storage::hummock::test_utils::gen_dummy_batch;
+use risingwave_storage::hummock::test_utils::*;
 
 use crate::test_utils::prepare_first_valid_version;
 

diff --git a/src/storage/hummock_test/src/hummock_storage_tests.rs b/src/storage/hummock_test/src/hummock_storage_tests.rs
@@ -41,7 +41,7 @@ use risingwave_meta::hummock::{CommitEpochInfo, NewTableFragmentInfo};
 use risingwave_rpc_client::HummockMetaClient;
 use risingwave_storage::hummock::local_version::pinned_version::PinnedVersion;
 use risingwave_storage::hummock::store::version::read_filter_for_version;
-use risingwave_storage::hummock::test_utils::StateStoreReadTestExt;
+use risingwave_storage::hummock::test_utils::*;
 use risingwave_storage::hummock::{CachePolicy, HummockStorage, LocalHummockStorage};
 use risingwave_storage::storage_value::StorageValue;
 use risingwave_storage::store::*;

diff --git a/src/storage/hummock_test/src/snapshot_tests.rs b/src/storage/hummock_test/src/snapshot_tests.rs
@@ -24,7 +24,7 @@ use risingwave_hummock_sdk::key::prefixed_range_with_vnode;
 use risingwave_hummock_sdk::HummockReadEpoch;
 use risingwave_meta::hummock::MockHummockMetaClient;
 use risingwave_rpc_client::HummockMetaClient;
-use risingwave_storage::hummock::test_utils::StateStoreReadTestExt;
+use risingwave_storage::hummock::test_utils::*;
 use risingwave_storage::hummock::{CachePolicy, HummockStorage};
 use risingwave_storage::storage_value::StorageValue;
 use risingwave_storage::store::{

diff --git a/src/storage/hummock_test/src/state_store_tests.rs b/src/storage/hummock_test/src/state_store_tests.rs
@@ -35,9 +35,7 @@ use risingwave_storage::hummock::iterator::change_log::test_utils::{
     apply_test_log_data, gen_test_data,
 };
 use risingwave_storage::hummock::iterator::test_utils::mock_sstable_store;
-use risingwave_storage::hummock::test_utils::{
-    count_stream, default_opts_for_test, StateStoreReadTestExt,
-};
+use risingwave_storage::hummock::test_utils::*;
 use risingwave_storage::hummock::{CachePolicy, HummockStorage};
 use risingwave_storage::memory::MemoryStateStore;
 use risingwave_storage::storage_value::StorageValue;

diff --git a/src/storage/hummock_test/src/sync_point_tests.rs b/src/storage/hummock_test/src/sync_point_tests.rs
@@ -34,7 +34,7 @@ use risingwave_rpc_client::HummockMetaClient;
 use risingwave_storage::compaction_catalog_manager::CompactionCatalogAgentRef;
 use risingwave_storage::hummock::compactor::compactor_runner::compact_with_agent;
 use risingwave_storage::hummock::compactor::CompactorContext;
-use risingwave_storage::hummock::test_utils::StateStoreReadTestExt;
+use risingwave_storage::hummock::test_utils::*;
 use risingwave_storage::hummock::{CachePolicy, GetObjectId, SstableObjectIdManager};
 use risingwave_storage::store::{LocalStateStore, NewLocalOptions, ReadOptions};
 use risingwave_storage::StateStore;

diff --git a/src/storage/hummock_test/src/test_utils.rs b/src/storage/hummock_test/src/test_utils.rs
@@ -39,7 +39,7 @@ use risingwave_storage::hummock::event_handler::HummockVersionUpdate;
 use risingwave_storage::hummock::iterator::test_utils::mock_sstable_store;
 use risingwave_storage::hummock::local_version::pinned_version::PinnedVersion;
 use risingwave_storage::hummock::observer_manager::HummockObserverNode;
-use risingwave_storage::hummock::test_utils::default_opts_for_test;
+use risingwave_storage::hummock::test_utils::*;
 use risingwave_storage::hummock::write_limiter::WriteLimiter;
 use risingwave_storage::hummock::HummockStorage;
 use risingwave_storage::storage_value::StorageValue;