Repository: risingwavelabs/risingwave
File: src/storage/hummock_trace/src/replay/worker.rs (Rust)

Code before:
    fn len(&self) -> usize {
        self.storages.len()
    }
}

#[cfg(test)]
mod tests {

    use std::ops::Bound;

    use bytes::Bytes;
    use mockall::predicate;

    use super::*;
    use crate::replay::{MockGlobalReplayInterface, MockLocalReplayInterface};
    use crate::{MockReplayIterStream, TracedBytes, TracedReadOptions};

    #[tokio::test]
    async fn test_handle_record() {
        let mut iters_map = HashMap::new();
        let mut local_storage_opts_map = HashMap::new();
        let mut local_storages = LocalStorages::new();
        let (res_tx, mut res_rx) = unbounded_channel();
        let get_table_id = 12;
        let iter_table_id = 14654;
        let get_fragment_id = 0;

        let read_options = TracedReadOptions::for_test(get_table_id);
        let iter_read_options = TracedReadOptions::for_test(iter_table_id);
        let op = Operation::get(Bytes::from(vec![123]), Some(123), read_options);

        let new_local_opts = TracedNewLocalOptions::for_test(get_table_id, get_fragment_id);

        let iter_local_opts = TracedNewLocalOptions::for_test(iter_table_id, get_fragment_id);
        let mut should_exit = false;
        let get_storage_type = StorageType::Local(0, 0);
        let record = Record::new(get_storage_type, 1, op);
        let mut mock_replay = MockGlobalReplayInterface::new();

        mock_replay.expect_new_local().times(1).returning(move |_| {
            let mut mock_local = MockLocalReplayInterface::new();

            mock_local
                .expect_get()
                .with(
                    predicate::eq(TracedBytes::from(vec![123])),
                    predicate::always(),
                )
                .returning(|_, _| Ok(Some(TracedBytes::from(vec![120]))));

            Box::new(mock_local)

Review comment (refactor):
Except your current modifications, the most important part you need to modify is to change file `src/storage/hummock_trace/src/opts.rs.`
```
In impl TracedNewLocalOptions {
pub(crate) fn for_test(table_id: u32, fragment_id: u32) -> Self
```
You need to set the fragment_id as 0 internally, and remove this parameter. Besides, check all for_test() in target branch and modify them to match the updated function signature.

Code after:
    fn len(&self) -> usize {
        self.storages.len()
    }
}

#[cfg(test)]
mod tests {

    use std::ops::Bound;

    use bytes::Bytes;
    use mockall::predicate;

    use super::*;
    use crate::replay::{MockGlobalReplayInterface, MockLocalReplayInterface};
    use crate::{MockReplayIterStream, TracedBytes, TracedReadOptions};

    #[tokio::test]
    async fn test_handle_record() {
        let mut iters_map = HashMap::new();
        let mut local_storage_opts_map = HashMap::new();
        let mut local_storages = LocalStorages::new();
        let (res_tx, mut res_rx) = unbounded_channel();
        let get_table_id = 12;
        let iter_table_id = 14654;

        let read_options = TracedReadOptions::for_test(get_table_id);
        let iter_read_options = TracedReadOptions::for_test(iter_table_id);
        let op = Operation::get(Bytes::from(vec![123]), Some(123), read_options);

        let new_local_opts = TracedNewLocalOptions::for_test(get_table_id);

        let iter_local_opts = TracedNewLocalOptions::for_test(iter_table_id);
        let mut should_exit = false;
        let get_storage_type = StorageType::Local(0, 0);
        let record = Record::new(get_storage_type, 1, op);
        let mut mock_replay = MockGlobalReplayInterface::new();

        mock_replay.expect_new_local().times(1).returning(move |_| {
            let mut mock_local = MockLocalReplayInterface::new();

            mock_local
                .expect_get()
                .with(
                    predicate::eq(TracedBytes::from(vec![123])),
                    predicate::always(),
                )
                .returning(|_, _| Ok(Some(TracedBytes::from(vec![120]))));

            Box::new(mock_local)
        });