use anyhow::{Ok, Result, anyhow};
use log::info;
use scopeguard::defer;
use tici::{
    config::{make_default_test_meta_config, make_default_test_worker_config},
    meta::{
        Bootstrap,
        metrics::{SHARD_COUNT_GAUGE_VEC, SHARD_EVENT_COUNTER_VEC},
        server::MetaServiceServer,
    },
};
use tici_common::index::IndexKey;
use tici_grpc::{ErrorCode, WorkerNodeHeartbeatRequest};
use tici_system::readable::ReadableDuration;
use tici_test_utils::test_utils::{test_suite::E2ETestSuite, *};

#[test]
fn test_meta_server_heartbeat() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap();
    defer!({
        fail::remove("worker_node_client_fail_fast");
    });
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        // Heartbeat to worker node
        info!("Start heartbeat test...");
        let config = make_default_test_meta_config(false);
        let (tx, rx) = tokio::sync::oneshot::channel::<()>();
        let config_clone = config.clone();
        Bootstrap::new_with_config(&config)
            .clear_all_meta(true)
            .await
            .unwrap();
        let mut server = MetaServiceServer::new(config_clone).await.unwrap();
        server.start(rx).await.unwrap();
        let mut client = MetaServiceServer::get_internal_client(
            config.server.get_advertise_addr(),
            &config.server.grpc_client,
        )
        .await;
        let req = WorkerNodeHeartbeatRequest::default();
        let mut stream = client
            .worker_node_heartbeat(req)
            .await
            .unwrap()
            .into_inner();
        let _ = stream.message().await.unwrap();
        server.stop(tx).await.unwrap();
        info!("Stop heartbeat test...");
    });
    Ok(())
}

#[test]
fn test_meta_server_with_worker() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap();
    defer!({
        fail::remove("worker_node_client_fail_fast");
    });
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
        defer!({
            fail::remove("worker_service_mock_meta_service_url");
        });
        info!("Start worker interaction test ...");
        let _v: Vec<u64> = vec![1; 102400000];

        let mut test_suite = E2ETestSuite::new_with_mock_cfg(
            make_default_test_meta_config(true),
            make_default_test_worker_config(true),
        )
        .await
        .unwrap();
        test_suite.wait_heartbeat().await;
        test_suite.stop_worker().await.unwrap(); // only stop workers

        assert_eq!(test_suite.collect_worker_nodes().await.len(), 1);
        let latest_status = test_suite
            .get_meta_service()
            .worker_node_mgr
            .get_worker_node(&test_suite.get_worker_node_template_addr())
            .await
            .unwrap()
            .get_status();

        {
            #[cfg(target_os = "linux")]
            assert!(latest_status.memory_usage > 0);

            #[cfg(not(target_os = "linux"))]
            {
                info!(
                    "Memory usage on non-Linux system={}",
                    latest_status.memory_usage
                );
            }
        }
    });
    Ok(())
}

#[test]
fn test_multiple_workers() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut test_suite = E2ETestSuite::new_with_mock_cfg(
            make_default_test_meta_config(true),
            make_default_test_worker_config(true),
        )
        .await
        .unwrap();

        let new_wn_addr = "127.0.0.1:8512".to_string();
        test_suite
            .add_mocked_worker_node(Some(new_wn_addr.clone()))
            .await;

        test_suite
            .must_create_index(&IndexKey::new_nullspace(1, 1001))
            .await
            .unwrap();
        test_suite
            .must_create_index(&IndexKey::new_nullspace(1, 1002))
            .await
            .unwrap();

        must_finally(async || {
            test_suite
                .get_worker_node_handle(&new_wn_addr)
                .worker_node_service
                .shard_count()
                == 1
        })
        .await;
        must_finally(async || {
            test_suite
                .get_worker_node_handle("127.0.0.1:8510")
                .worker_node_service
                .shard_count()
                == 1
        })
        .await;
        must_finally(async || {
            test_suite
                .get_worker_node_handle("127.0.0.1:8512")
                .worker_node_service
                .shard_count()
                == 1
        })
        .await;

        test_suite
            .must_drop_index(&IndexKey::new_nullspace(1, 1001))
            .await
            .unwrap();
        test_suite
            .must_drop_index(&IndexKey::new_nullspace(1, 1002))
            .await
            .unwrap();

        test_suite.stop_worker().await.unwrap();
        assert_eq!(test_suite.collect_shards_from_restore().await.len(), 0);
        assert_eq!(test_suite.collect_indexes_from_restore().await.len(), 0);
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_advertise_addr() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
        defer!({
            fail::remove("worker_service_mock_meta_service_url");
        });
        let mut wn_config = make_default_test_worker_config(true);
        wn_config.server.addr = "127.0.0.1:8510".to_string();
        wn_config
            .server
            .set_advertise_addr("127.0.0.1:8510".to_string());
        assert!(!wn_config.security.is_valid());
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(make_default_test_meta_config(true), wn_config)
                .await
                .unwrap();

        assert!(
            !test_suite
                .get_worker_node_handle("127.0.0.1:8510")
                .shard_write_manager
                .as_mock()
                .skip_heartbeat
                .load(std::sync::atomic::Ordering::Relaxed)
        );

        test_suite
            .must_create_index(&IndexKey::new_nullspace(1, 1001))
            .await
            .unwrap();

        must_finally(async || {
            test_suite
                .get_worker_node_handle("127.0.0.1:8510")
                .worker_node_service
                .shard_count()
                == 1
        })
        .await;

        must_finally(async || {
            test_suite
                .get_meta_service()
                .worker_node_mgr
                .get_worker_node("127.0.0.1:8510")
                .await
                .is_some()
        })
        .await;

        test_suite
            .must_drop_index(&IndexKey::new_nullspace(1, 1001))
            .await
            .unwrap();
        test_suite.stop_worker().await.unwrap();
        assert_eq!(test_suite.collect_shards_from_restore().await.len(), 0);
        assert_eq!(test_suite.collect_indexes_from_restore().await.len(), 0);
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_resche_shards_on_two_wns() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut config = make_default_test_meta_config(false);
    config.worker_node.worker_node_down_timeout = ReadableDuration::secs(1);
    config.worker_node.health_check_interval = ReadableDuration::secs(1);

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        let new_wn_addr = "127.0.0.1:8512".to_string();
        test_suite
            .add_mocked_worker_node(Some(new_wn_addr.clone()))
            .await;

        test_suite
            .must_create_index(&IndexKey::new_nullspace(1, 1001))
            .await
            .unwrap();
        test_suite
            .must_create_index(&IndexKey::new_nullspace(1, 1002))
            .await
            .unwrap();

        must_finally(async || {
            test_suite
                .get_worker_node_handle("127.0.0.1:8512")
                .worker_node_service
                .shard_count()
                == 1
        })
        .await;

        must_finally(async || {
            test_suite
                .get_worker_node_handle("127.0.0.1:8510")
                .worker_node_service
                .shard_count()
                == 1
        })
        .await;

        test_suite
            .get_worker_node_handle("127.0.0.1:8512")
            .shard_write_manager
            .as_mock()
            .skip_heartbeat
            .store(true, std::sync::atomic::Ordering::Relaxed);

        tokio::time::sleep(std::time::Duration::from_secs(2)).await;
        assert_eq!(
            test_suite
                .get_meta_service()
                .worker_node_mgr
                .get_worker_nodes_len()
                .await,
            1
        );
        let shards = test_suite.collect_shards_from_restore().await;
        for shard in shards.iter() {
            assert_eq!(shard.writer, "127.0.0.1:8510".to_string());
        }

        tokio::time::sleep(std::time::Duration::from_secs(2)).await;

        test_suite.stop_worker().await.unwrap();
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_resche_shards_on_one_wn() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut config = make_default_test_meta_config(false);
    config.worker_node.worker_node_down_timeout = ReadableDuration::secs(1);
    config.worker_node.health_check_interval = ReadableDuration::secs(1);

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await.unwrap();

        must_finally(async || {
            test_suite
                .get_worker_node_handle("127.0.0.1:8510")
                .worker_node_service
                .shard_count()
                == 1
        })
        .await;

        let shard = test_suite.get_first_shard().await;

        test_suite
            .get_meta_service()
            .shard_task_manager
            .on_msg_manual_reschedule_by_shard(
                &index_key,
                shard.range.encoded_start(),
                shard.range.encoded_end(),
                None,
            )
            .await
            .unwrap();

        assert_eq!(
            test_suite
                .get_meta_service()
                .worker_node_mgr
                .get_worker_nodes_len()
                .await,
            1
        );
        assert_eq!(
            test_suite
                .collect_shards_by_worker_node(&"127.0.0.1:8510".to_string())
                .await
                .len(),
            1
        );
        assert_eq!(
            test_suite
                .get_meta_service()
                .worker_node_mgr
                .get_worker_node("127.0.0.1:8510")
                .await
                .unwrap()
                .shards_len()
                .await,
            1
        );

        let shards = test_suite.collect_shards_from_restore().await;
        for shard in shards.iter() {
            assert_eq!(shard.writer, "127.0.0.1:8510".to_string());
        }

        test_suite
            .get_worker_node_handle("127.0.0.1:8510")
            .shard_write_manager
            .as_mock()
            .skip_heartbeat
            .store(true, std::sync::atomic::Ordering::Relaxed);

        tokio::time::sleep(std::time::Duration::from_secs(2)).await;

        assert_eq!(
            test_suite
                .get_meta_service()
                .worker_node_mgr
                .get_worker_nodes_len()
                .await,
            0
        );

        test_suite.stop_worker().await.unwrap();
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_manual_reschedule_removes_old_worker_schedule() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").map_err(anyhow::Error::msg)?;
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut config = make_default_test_meta_config(false);
    config.worker_node.worker_node_down_timeout = ReadableDuration::secs(3600);
    config.worker_node.health_check_interval = ReadableDuration::secs(3600);

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true)).await?;

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await?;

        let shard_before = test_suite.get_first_shard().await;
        assert_eq!(shard_before.writer, "127.0.0.1:8510".to_string());

        let new_wn_addr = "127.0.0.1:8512".to_string();
        test_suite
            .add_mocked_worker_node(Some(new_wn_addr.clone()))
            .await;

        let (code, err_msg) = test_suite
            .get_meta_service()
            .shard_task_manager
            .on_msg_manual_reschedule_by_shard(
                &index_key,
                shard_before.range.encoded_start(),
                shard_before.range.encoded_end(),
                Some(new_wn_addr.clone()),
            )
            .await?;
        assert_eq!(
            code,
            tici_grpc::ErrorCode::Success,
            "manual_reschedule failed: {}",
            err_msg
        );

        let shard_after = test_suite
            .collect_shards_from_restore()
            .await
            .first()
            .cloned()
            .ok_or_else(|| anyhow!("missing shard after manual reschedule"))?;
        assert_eq!(shard_after.writer, new_wn_addr);

        let old_wn = test_suite
            .get_meta_service()
            .worker_node_mgr
            .get_worker_node("127.0.0.1:8510")
            .await
            .ok_or_else(|| anyhow!("missing old worker node in meta"))?;
        let new_wn = test_suite
            .get_meta_service()
            .worker_node_mgr
            .get_worker_node("127.0.0.1:8512")
            .await
            .ok_or_else(|| anyhow!("missing new worker node in meta"))?;
        assert!(
            old_wn
                .get_shard_schedule(&shard_before.shard_id)
                .await
                .is_none()
        );
        assert!(
            new_wn
                .get_shard_schedule(&shard_before.shard_id)
                .await
                .is_some()
        );

        test_suite.stop_worker().await?;
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await?;
        Ok(())
    })?;

    Ok(())
}

#[test]
fn test_patrol_reschedules_shard_when_worker_omits_it_from_heartbeat() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").map_err(anyhow::Error::msg)?;
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut cfg = make_default_test_meta_config(false);
    cfg.shard.enable_patrol();
    cfg.shard.shard_no_heartbeat_tick = 1;
    // Keep the patrol runner's background dead-shard scan out of this test so we
    // can deterministically trigger exactly one patrol round with
    // `debug_run_one`.
    cfg.shard.check_dead_shard_tick = u64::MAX;
    cfg.worker_node.worker_node_down_timeout = ReadableDuration::secs(3600);
    cfg.worker_node.health_check_interval = ReadableDuration::secs(3600);
    let shard_heartbeat_timeout_wait_ms =
        cfg.shard.patrol_interval.as_millis() * (cfg.shard.shard_no_heartbeat_tick + 1);

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(cfg, make_default_test_worker_config(true)).await?;
        let new_wn_addr = "127.0.0.1:8512".to_string();
        test_suite
            .add_mocked_worker_node(Some(new_wn_addr.clone()))
            .await;

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await?;

        let shard_before = test_suite.get_first_shard().await;
        let shard_id = shard_before.shard_id;
        let old_writer = shard_before.writer.clone();
        let rescheduled_writer = if old_writer == new_wn_addr {
            "127.0.0.1:8510".to_string()
        } else {
            new_wn_addr.clone()
        };

        test_suite
            .get_worker_node_handle(&old_writer)
            .shard_write_manager
            .as_mock()
            .mock_skip_shard_heartbeat(shard_id);

        tokio::time::sleep(std::time::Duration::from_millis(
            shard_heartbeat_timeout_wait_ms,
        ))
        .await;

        assert_eq!(
            test_suite
                .get_meta_service()
                .shard_task_manager
                .get_shard_patrol()
                .debug_run_one(0)
                .await?
                .sched_no_heartbeat_count,
            1
        );

        must_finally(async || {
            if let std::result::Result::Ok(shard) =
                test_suite.get_shard_from_restore(shard_id).await
            {
                shard.writer == rescheduled_writer
            } else {
                false
            }
        })
        .await;

        must_finally(async || {
            let new_worker = test_suite.get_worker_node_handle(&rescheduled_writer);
            test_suite
                .get_meta_service()
                .worker_node_mgr
                .get_worker_node(&old_writer)
                .await
                .unwrap()
                .get_shard_schedule(&shard_id)
                .await
                .is_none()
                && new_worker.worker_node_service.contains_shard(shard_id)
                && test_suite
                    .get_meta_service()
                    .worker_node_mgr
                    .get_worker_node(&rescheduled_writer)
                    .await
                    .unwrap()
                    .get_shard_schedule(&shard_id)
                    .await
                    .is_some()
        })
        .await;

        test_suite.stop_worker().await?;
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await?;
        Ok(())
    })?;

    Ok(())
}

#[test]
fn test_with_no_worker() -> Result<()> {
    // Worker node timeouts first, then a shard timeouts in patrol loop.
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut cfg = make_default_test_meta_config(true);
    cfg.shard.shard_no_heartbeat_tick = 1;
    cfg.shard.enable_patrol();
    cfg.worker_node.worker_node_down_timeout = ReadableDuration::secs(6);
    cfg.worker_node.health_check_interval = ReadableDuration::millis(500);
    let wn_cfg = make_default_test_worker_config(true);
    rt.block_on(async {
        let suite = E2ETestSuite::new_with_mock_cfg(cfg.clone(), wn_cfg.clone())
            .await
            .unwrap();
        let index_key = IndexKey::new_nullspace(1, 1001);
        suite.must_create_index(&index_key).await.unwrap();

        suite.validate_add_index_result(&index_key).await;

        suite
            .get_worker_node_handle("127.0.0.1:8510")
            .shard_write_manager
            .as_mock()
            .skip_heartbeat
            .store(true, std::sync::atomic::Ordering::Relaxed);

        tokio::time::sleep(std::time::Duration::from_millis(
            cfg.shard.patrol_interval.as_millis() * cfg.shard.shard_no_heartbeat_tick * 3,
        ))
        .await;
        assert_eq!(
            suite
                .get_meta_service()
                .shard_task_manager
                .get_shard_patrol()
                .debug_run_one(0)
                .await
                .unwrap()
                .sched_no_heartbeat_count,
            1
        );
        assert_eq!(
            SHARD_COUNT_GAUGE_VEC
                .get_metric_with_label_values(&["dead"])
                .iter()
                .next()
                .unwrap()
                .get(),
            1
        );

        let shard = suite
            .collect_shards_from_restore()
            .await
            .first()
            .unwrap()
            .clone();
        assert!(!shard.writer.is_empty());
        tokio::time::sleep(std::time::Duration::from_millis(
            cfg.shard.patrol_interval.as_millis(),
        ))
        .await;
        assert!(
            SHARD_EVENT_COUNTER_VEC
                .get_metric_with_label_values(&["resche_ok"])
                .iter()
                .next()
                .unwrap()
                .get()
                > 0
        );
        assert_eq!(
            SHARD_EVENT_COUNTER_VEC
                .get_metric_with_label_values(&["resche_failed"])
                .iter()
                .next()
                .unwrap()
                .get(),
            0
        );
        tokio::time::sleep(std::time::Duration::from_millis(
            cfg.worker_node.worker_node_down_timeout.as_millis(),
        ))
        .await;

        // Now, the only worker node is considered down.
        assert_eq!(
            suite
                .collect_shards_from_restore()
                .await
                .first()
                .unwrap()
                .writer,
            ""
        );

        // Now, we recover the node
        suite
            .get_worker_node_handle("127.0.0.1:8510")
            .shard_write_manager
            .as_mock()
            .skip_heartbeat
            .store(false, std::sync::atomic::Ordering::Relaxed);

        tokio::time::sleep(std::time::Duration::from_millis(
            cfg.shard.patrol_interval.as_millis() * cfg.shard.shard_no_heartbeat_tick * 3,
        ))
        .await;

        assert_eq!(
            suite
                .get_meta_service()
                .worker_node_mgr
                .get_worker_node("127.0.0.1:8510")
                .await
                .unwrap()
                .shards_len()
                .await,
            1
        );

        assert_eq!(
            suite
                .get_worker_node_handle("127.0.0.1:8510")
                .worker_node_service
                .shard_count(),
            1
        );

        assert_ne!(
            suite
                .collect_shards_from_restore()
                .await
                .first()
                .unwrap()
                .writer,
            ""
        );
    });
    Ok(())
}

#[test]
fn test_progress_unwind_is_recovered_by_patrol_after_meta_restart() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut cfg = make_default_test_meta_config(false);
    cfg.shard.enable_patrol();
    cfg.shard.shard_no_heartbeat_tick = 3;
    cfg.worker_node.worker_node_down_timeout = ReadableDuration::secs(3600);
    cfg.worker_node.health_check_interval = ReadableDuration::secs(3600);
    let restart_cfg = cfg.clone();

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(cfg, make_default_test_worker_config(true))
                .await
                .unwrap();
        test_suite
            .add_mocked_worker_node(Some("127.0.0.1:8512".to_string()))
            .await;

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await.unwrap();

        let shard = test_suite.get_first_shard().await;
        let shard_id = shard.shard_id;
        let original_seq = shard.seq;
        let old_writer = shard.writer.clone();
        let append_cdc_file = create_cdc_file_name(2);

        let (_, append_resp) = test_suite
            .do_append_delta_frag(
                old_writer.as_str(),
                shard_id,
                &index_key,
                append_cdc_file.clone(),
                vec![],
            )
            .await
            .unwrap();
        assert_eq!(append_resp.status, ErrorCode::Success as i32);
        let shard_after_empty_append = test_suite.get_shard_from_restore(shard_id).await.unwrap();
        assert_eq!(shard_after_empty_append.seq, original_seq);
        assert_eq!(
            shard_after_empty_append.progress.cdc_s3_last_file,
            append_cdc_file
        );

        must_finally(async || {
            test_suite
                .check_get_shard_manifest(&index_key, shard_id, original_seq, vec![])
                .await
        })
        .await;

        let (_, rewind_resp) = test_suite
            .do_append_delta_frag(
                old_writer.as_str(),
                shard_id,
                &index_key,
                create_cdc_file_name(1),
                vec![],
            )
            .await
            .unwrap();
        assert_eq!(rewind_resp.status, ErrorCode::ProgressUnwinded as i32);

        must_finally(async || {
            if let std::result::Result::Ok(shard) =
                test_suite.get_shard_from_restore(shard_id).await
            {
                shard.writer.is_empty()
            } else {
                false
            }
        })
        .await;

        must_finally(async || {
            let worker = test_suite.get_worker_node_handle(&old_writer);
            !worker.worker_node_service.contains_shard(shard_id)
        })
        .await;

        test_suite.restart_meta(restart_cfg.clone()).await.unwrap();
        test_suite.wait_heartbeat().await;

        must_finally(async || {
            if let std::result::Result::Ok(shard) =
                test_suite.get_shard_from_restore(shard_id).await
            {
                !shard.writer.is_empty()
            } else {
                false
            }
        })
        .await;

        let recovered = test_suite.get_shard_from_restore(shard_id).await.unwrap();
        let recovered_writer = recovered.writer.clone();
        assert!(
            !recovered_writer.is_empty(),
            "patrol should reassign the shard after progress unwind"
        );
        must_finally(async || {
            test_suite
                .get_worker_node_handle(&recovered_writer)
                .worker_node_service
                .contains_shard(shard_id)
        })
        .await;

        test_suite.stop_worker().await.unwrap();
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}