Repository: tikv/tikv
File: tests/integrations/import/test_sst_service.rs (Rust)

Code before:
    let ret = block_on(async move {
        let mut ret: Vec<KvPair> = vec![];
        while let Some(resp) = stream.next().await {
            match resp {
                Ok(mut resp) => {
                    if resp.has_key_error() || resp.has_region_error() {
                        break;
                    }
                    let pairs = resp.take_pairs();
                    ret.append(&mut pairs.into());
                }
                Err(e) => {
                    println!("receive error: {:?}", e);
                    break;
                }
            }
        }
        ret
    });
    assert_eq!(ret.len(), (data_count - 1000) as usize * 4);
    req.set_mode(SwitchMode::Normal);
    import.switch_mode(&req).unwrap();
}

#[test]
fn test_duplicate_detect_with_client_stop() {
    let (_cluster, ctx, _, import) = new_cluster_and_tikv_import_client();
    let mut req = SwitchModeRequest::default();
    req.set_mode(SwitchMode::Import);
    import.switch_mode(&req).unwrap();

    let data_count: u64 = 4096;
    for commit_ts in 0..4 {
        let mut meta = new_sst_meta(0, 0);
        meta.set_region_id(ctx.get_region_id());
        meta.set_region_epoch(ctx.get_region_epoch().clone());

        let mut keys = vec![];
        let mut values = vec![];
        for i in 1000..data_count {
            let key = i.to_string();
            keys.push(key.as_bytes().to_vec());
            values.push(key.as_bytes().to_vec());
        }
        let resp = send_write_sst(&import, &meta, keys, values, commit_ts).unwrap();
        for m in resp.metas.into_iter() {
            must_ingest_sst(&import, ctx.clone(), m.clone());
        }
    }

    let mut duplicate = DuplicateDetectRequest::default();

Review comment (bug):
I guess test cases with failpoints enabled should be put in the `failpoints` dir, or the failpoint may race with other cases...

Code after:
    let ret = block_on(async move {
        let mut ret: Vec<KvPair> = vec![];
        while let Some(resp) = stream.next().await {
            match resp {
                Ok(mut resp) => {
                    if resp.has_key_error() || resp.has_region_error() {
                        break;
                    }
                    let pairs = resp.take_pairs();
                    ret.append(&mut pairs.into());
                }
                Err(e) => {
                    println!("receive error: {:?}", e);
                    break;
                }
            }
        }
        ret
    });
    assert_eq!(ret.len(), (data_count - 1000) as usize * 4);
    req.set_mode(SwitchMode::Normal);
    import.switch_mode(&req).unwrap();
}

#[test]
fn test_suspend_import() {
    let (_cluster, ctx, tikv, import) = new_cluster_and_tikv_import_client();
    let sst_range = (0, 10);
    let write = |sst_range: (u8, u8)| {
        let mut meta = new_sst_meta(0, 0);
        meta.set_region_id(ctx.get_region_id());
        meta.set_region_epoch(ctx.get_region_epoch().clone());

        let mut keys = vec![];
        let mut values = vec![];
        for i in sst_range.0..sst_range.1 {
            keys.push(vec![i]);
            values.push(vec![i]);
        }
        send_write_sst(&import, &meta, keys, values, 1)
    };
    let ingest = |sst_meta: &SstMeta| {
        let mut ingest = IngestRequest::default();
        ingest.set_context(ctx.clone());
        ingest.set_sst(sst_meta.clone());
        import.ingest(&ingest)
    };
    let multi_ingest = |sst_metas: &[SstMeta]| {
        let mut multi_ingest = MultiIngestRequest::default();
        multi_ingest.set_context(ctx.clone());
        multi_ingest.set_ssts(sst_metas.to_vec().into());