// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

#[cfg(any(test, feature = "testexport"))]
use std::sync::OnceLock;
use std::{
    collections::{
        BTreeMap,
        Bound::{self, Excluded, Included, Unbounded},
    },
    default::Default,
    fmt::{self, Debug, Display, Formatter},
    ops::RangeBounds,
    sync::{
        Arc, RwLock,
        atomic::{AtomicBool, Ordering},
    },
};

use collections::HashMap;
use engine_traits::{CF_DEFAULT, CF_LOCK, CF_WRITE, CfName, IterOptions, ReadOptions};
use futures::{Future, Stream, stream};
use kvproto::kvrpcpb::{ApiVersion, Context};
#[cfg(any(test, feature = "testexport"))]
use txn_types::LockType;
use txn_types::{Key, Value, WriteRef, WriteType};

use super::SnapContext;
use crate::{
    DummySnapshotExt, Engine, Error as EngineError, ErrorInner as EngineErrorInner, Iterator,
    Modify, OnAppliedCb, Result as EngineResult, Snapshot, TrackerToken, WriteData, WriteEvent,
};

type RwLockTree = RwLock<BTreeMap<Key, Value>>;

/// The `BTreeEngine` is an in-memory engine used by logic tests and benchmarks.
/// It deliberately provides no persistence or physical-SST behavior.
#[derive(Clone)]
pub struct BTreeEngine {
    cf_names: Vec<CfName>,
    cf_contents: Vec<Arc<RwLockTree>>,
    raft_extension: crate::FakeExtension,
    not_leader: Arc<AtomicBool>,
    api_version: ApiVersion,
}

impl BTreeEngine {
    pub fn new(cfs: &[CfName]) -> Self {
        Self::new_with_api_version(cfs, ApiVersion::V1)
    }

    /// Creates an in-memory engine whose projected kvengine snapshot uses the
    /// requested API key layout.
    pub fn new_with_api_version(cfs: &[CfName], api_version: ApiVersion) -> Self {
        let mut cf_names = vec![];
        let mut cf_contents = vec![];

        // create default cf if missing
        if !cfs.contains(&CF_DEFAULT) {
            cf_names.push(CF_DEFAULT);
            cf_contents.push(Arc::new(RwLock::new(BTreeMap::new())))
        }

        for cf in cfs.iter() {
            cf_names.push(*cf);
            cf_contents.push(Arc::new(RwLock::new(BTreeMap::new())))
        }

        Self {
            cf_names,
            cf_contents,
            raft_extension: crate::FakeExtension,
            not_leader: Arc::new(AtomicBool::new(false)),
            api_version,
        }
    }

    pub fn get_cf(&self, cf: CfName) -> Arc<RwLockTree> {
        let index = self
            .cf_names
            .iter()
            .position(|&c| c == cf)
            .expect("CF not exist!");
        self.cf_contents[index].clone()
    }

    /// Makes subsequent request prechecks return a not-leader error. This is a
    /// lightweight error-injection hook used by request-routing tests.
    pub fn trigger_not_leader(&self) {
        self.not_leader.store(true, Ordering::SeqCst);
    }

    fn not_leader_error(&self) -> EngineError {
        let mut error = kvproto::errorpb::Error::default();
        error.mut_not_leader().set_region_id(100);
        EngineError::from(EngineErrorInner::Request(error))
    }
}

impl Default for BTreeEngine {
    fn default() -> Self {
        let cfs = &[CF_WRITE, CF_DEFAULT, CF_LOCK];
        Self::new(cfs)
    }
}

impl Engine for BTreeEngine {
    type Snap = BTreeEngineSnapshot;
    type Local = BTreeEngine;

    fn kv_engine(&self) -> Option<BTreeEngine> {
        Some(self.clone())
    }

    type RaftExtension = crate::FakeExtension;

    fn raft_extension(&self) -> &Self::RaftExtension {
        &self.raft_extension
    }

    fn modify_on_kv_engine(&self, region_modifies: HashMap<u64, Vec<Modify>>) -> EngineResult<()> {
        write_modifies(self, region_modifies.into_values().flatten().collect())
    }

    fn precheck_write_with_ctx(&self, _ctx: &Context) -> EngineResult<()> {
        if !self.not_leader.load(Ordering::SeqCst) {
            return Ok(());
        }
        Err(self.not_leader_error())
    }

    type WriteRes = impl Stream<Item = WriteEvent> + Send;
    fn async_write(
        &self,
        _ctx: &Context,
        batch: WriteData,
        subscribed: u8,
        on_applied: Option<OnAppliedCb>,
        _tracker: Option<TrackerToken>,
    ) -> Self::WriteRes {
        let is_empty = batch.modifies.is_empty();
        let mut result = if is_empty {
            Err(EngineError::from(EngineErrorInner::EmptyRequest))
        } else {
            write_modifies(self, batch.modifies)
        };
        if let Some(callback) = on_applied {
            callback(&mut result);
        }

        let mut events = Vec::with_capacity(WriteEvent::event_capacity(subscribed));
        if !is_empty && WriteEvent::subscribed_proposed(subscribed) {
            events.push(WriteEvent::Proposed);
        }
        if !is_empty && WriteEvent::subscribed_committed(subscribed) {
            events.push(WriteEvent::Committed);
        }
        events.push(WriteEvent::Finished(result));
        stream::iter(events)
    }

    type SnapshotRes = impl Future<Output = EngineResult<Self::Snap>> + Send;
    fn async_snapshot(&mut self, _ctx: SnapContext<'_>) -> Self::SnapshotRes {
        let result = if self.not_leader.load(Ordering::SeqCst) {
            Err(self.not_leader_error())
        } else {
            Ok(BTreeEngineSnapshot::new(self))
        };
        futures::future::ready(result)
    }
}

impl Display for BTreeEngine {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "BTreeEngine",)
    }
}

impl Debug for BTreeEngine {
    // TODO: Provide more debug info.
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "BTreeEngine",)
    }
}

pub struct BTreeEngineIterator {
    tree: Arc<RwLockTree>,
    cur_key: Option<Key>,
    cur_value: Option<Value>,
    valid: bool,
    bounds: (Bound<Key>, Bound<Key>),
}

impl BTreeEngineIterator {
    pub fn new(tree: Arc<RwLockTree>, iter_opt: IterOptions) -> BTreeEngineIterator {
        let lower_bound = match iter_opt.lower_bound() {
            None => Unbounded,
            Some(key) => Included(Key::from_raw(key)),
        };

        let upper_bound = match iter_opt.upper_bound() {
            None => Unbounded,
            Some(key) => Excluded(Key::from_raw(key)),
        };
        let bounds = (lower_bound, upper_bound);
        Self {
            tree,
            cur_key: None,
            cur_value: None,
            valid: false,
            bounds,
        }
    }

    /// In general, there are 2 endpoints in a range, the left one and the right
    /// one. This method will seek to the left one if left is `true`, else seek
    /// to the right one. Returns true when the endpoint is valid, which means
    /// the endpoint exist and in `self.bounds`.
    fn seek_to_range_endpoint(&mut self, range: (Bound<Key>, Bound<Key>), left: bool) -> bool {
        let tree = self.tree.read().unwrap();
        let mut range = tree.range(range);
        let item = if left {
            range.next() // move to the left endpoint
        } else {
            range.next_back() // move to the right endpoint
        };
        match item {
            Some((k, v)) if self.bounds.contains(k) => {
                self.cur_key = Some(k.clone());
                self.cur_value = Some(v.clone());
                self.valid = true;
            }
            _ => {
                self.valid = false;
            }
        }
        self.valid().unwrap()
    }
}

impl Iterator for BTreeEngineIterator {
    fn next(&mut self) -> EngineResult<bool> {
        let range = (Excluded(self.cur_key.clone().unwrap()), Unbounded);
        Ok(self.seek_to_range_endpoint(range, true))
    }

    fn prev(&mut self) -> EngineResult<bool> {
        let range = (Unbounded, Excluded(self.cur_key.clone().unwrap()));
        Ok(self.seek_to_range_endpoint(range, false))
    }

    fn seek(&mut self, key: &Key) -> EngineResult<bool> {
        let range = (Included(key.clone()), Unbounded);
        Ok(self.seek_to_range_endpoint(range, true))
    }

    fn seek_for_prev(&mut self, key: &Key) -> EngineResult<bool> {
        let range = (Unbounded, Included(key.clone()));
        Ok(self.seek_to_range_endpoint(range, false))
    }

    fn seek_to_first(&mut self) -> EngineResult<bool> {
        let range = (self.bounds.0.clone(), self.bounds.1.clone());
        Ok(self.seek_to_range_endpoint(range, true))
    }

    fn seek_to_last(&mut self) -> EngineResult<bool> {
        let range = (self.bounds.0.clone(), self.bounds.1.clone());
        Ok(self.seek_to_range_endpoint(range, false))
    }

    #[inline]
    fn valid(&self) -> EngineResult<bool> {
        Ok(self.valid)
    }

    fn key(&self) -> &[u8] {
        assert!(self.valid().unwrap());
        self.cur_key.as_ref().unwrap().as_encoded()
    }

    fn value(&self) -> &[u8] {
        assert!(self.valid().unwrap());
        self.cur_value.as_ref().unwrap().as_slice()
    }
}

impl Snapshot for BTreeEngineSnapshot {
    type Iter = BTreeEngineIterator;
    type Ext<'a> = DummySnapshotExt;

    fn get(&self, key: &Key) -> EngineResult<Option<Value>> {
        self.get_cf(CF_DEFAULT, key)
    }

    fn get_cf(&self, cf: CfName, key: &Key) -> EngineResult<Option<Value>> {
        let tree_cf = self.inner_engine.get_cf(cf);
        let tree = tree_cf.read().unwrap();
        let v = tree.get(key);
        match v {
            None => Ok(None),
            Some(v) => Ok(Some(v.clone())),
        }
    }

    fn get_cf_opt(&self, _: ReadOptions, cf: CfName, key: &Key) -> EngineResult<Option<Value>> {
        self.get_cf(cf, key)
    }

    #[inline]
    fn iter(&self, cf: CfName, iter_opt: IterOptions) -> EngineResult<Self::Iter> {
        let tree = self.inner_engine.get_cf(cf);
        Ok(BTreeEngineIterator::new(tree, iter_opt))
    }

    fn ext(&self) -> DummySnapshotExt {
        DummySnapshotExt
    }

    #[cfg(any(test, feature = "testexport"))]
    fn get_kvengine_snap(&self) -> Option<&kvengine::SnapAccess> {
        Some(&self.kvengine.get_or_init(|| self.build_kvengine()).snap)
    }
}

#[derive(Debug, Clone)]
pub struct BTreeEngineSnapshot {
    inner_engine: Arc<BTreeEngine>,
    #[cfg(any(test, feature = "testexport"))]
    kvengine: Arc<OnceLock<KvEngineBacking>>,
}

#[cfg(any(test, feature = "testexport"))]
struct KvEngineBacking {
    // Fields are dropped in declaration order. Keep the snapshot before its
    // owning engine so no snapshot outlives the engine during teardown.
    snap: kvengine::SnapAccess,
    _engine: kvengine::test_engine::TestEngine,
}

#[cfg(any(test, feature = "testexport"))]
impl Debug for KvEngineBacking {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("KvEngineBacking").finish_non_exhaustive()
    }
}

impl BTreeEngineSnapshot {
    pub fn new(engine: &BTreeEngine) -> Self {
        // Writes acquire every CF lock before changing any of them. Reading the
        // locks in the same order therefore produces one cross-CF point-in-time
        // view rather than a collection affected by later writes.
        let trees = engine
            .cf_contents
            .iter()
            .map(|tree| tree.read().unwrap())
            .collect::<Vec<_>>();
        let cf_contents = trees
            .iter()
            .map(|tree| Arc::new(RwLock::new((*tree).clone())))
            .collect();
        Self {
            inner_engine: Arc::new(BTreeEngine {
                cf_names: engine.cf_names.clone(),
                cf_contents,
                raft_extension: crate::FakeExtension,
                not_leader: engine.not_leader.clone(),
                api_version: engine.api_version,
            }),
            #[cfg(any(test, feature = "testexport"))]
            kvengine: Arc::new(OnceLock::new()),
        }
    }

    /// Projects the generic test state into the current production MVCC
    /// layout. Most BTree tests never call this: it is built lazily only for
    /// code paths which explicitly request a `kvengine::SnapAccess`.
    #[cfg(any(test, feature = "testexport"))]
    fn build_kvengine(&self) -> KvEngineBacking {
        use kvengine::{
            EXTRA_CF, LOCK_CF, UserMeta, WRITE_CF, WriteBatch,
            test_engine::{ApplyTask, new_test_engine, new_test_engine_api_v2},
        };

        let (engine, apply_tx) = match self.inner_engine.api_version {
            ApiVersion::V1 | ApiVersion::V1ttl => new_test_engine(),
            ApiVersion::V2 => new_test_engine_api_v2(),
            ApiVersion::V3 => {
                unreachable!("API V3 is not supported on the cloud-engine branch")
            }
        };
        let mut batch = WriteBatch::default();
        batch.reset(1);

        let default_values = self
            .inner_engine
            .get_cf(CF_DEFAULT)
            .read()
            .unwrap()
            .iter()
            .filter_map(|(key, value)| {
                let start_ts = key.decode_ts().ok()?;
                let raw_key = key.clone().truncate_ts().ok()?.into_raw().ok()?;
                Some(((raw_key, start_ts.into_inner()), value.clone()))
            })
            .collect::<HashMap<_, _>>();

        for (key, value) in self.inner_engine.get_cf(CF_LOCK).read().unwrap().iter() {
            let raw_key = key.to_raw().unwrap();
            let value = match txn_types::parse_lock(value).unwrap() {
                tikv_util::Either::Left(mut lock) => {
                    if lock.lock_type == LockType::Put && lock.short_value.is_none() {
                        lock.short_value = default_values
                            .get(&(raw_key.clone(), lock.ts.into_inner()))
                            .cloned();
                    }
                    lock.to_bytes()
                }
                tikv_util::Either::Right(shared_locks) => shared_locks.to_bytes(),
            };
            batch.put(LOCK_CF, &raw_key, &value, 0, &[], 0);
        }

        for (key, value) in self.inner_engine.get_cf(CF_WRITE).read().unwrap().iter() {
            let commit_ts = key.decode_ts().unwrap().into_inner();
            let raw_key = key.clone().truncate_ts().unwrap().into_raw().unwrap();
            let write = WriteRef::parse(value).unwrap();
            let start_ts = write.start_ts.into_inner();
            let user_meta = UserMeta::new(start_ts, commit_ts).to_array();
            match write.write_type {
                WriteType::Put => {
                    let value = write.short_value.map_or_else(
                        || {
                            default_values
                                .get(&(raw_key.clone(), start_ts))
                                .cloned()
                                .expect("put write must have a value")
                        },
                        <[u8]>::to_vec,
                    );
                    batch.put(WRITE_CF, &raw_key, &value, 0, &user_meta, commit_ts);
                }
                WriteType::Delete => {
                    batch.put(WRITE_CF, &raw_key, &[], 0, &user_meta, commit_ts);
                }
                WriteType::Lock => {
                    let extra_key = kvengine::encode_extra_txn_status_key(&raw_key, start_ts);
                    batch.put(EXTRA_CF, &extra_key, &[0], 0, &user_meta, commit_ts);
                }
                WriteType::Rollback => {
                    let extra_key = kvengine::encode_extra_txn_status_key(&raw_key, start_ts);
                    let rollback_meta = UserMeta::new(start_ts, 0).to_array();
                    batch.put(EXTRA_CF, &extra_key, &[0], 0, &rollback_meta, start_ts);
                }
            }
        }

        let (result_tx, result_rx) = tikv_util::mpsc::bounded(1);
        apply_tx.send(ApplyTask::new_wb(batch, result_tx)).unwrap();
        result_rx.recv().unwrap().unwrap();
        let snap = engine.get_shard(1).unwrap().new_snap_access();
        KvEngineBacking {
            snap,
            _engine: engine,
        }
    }
}

fn write_modifies(engine: &BTreeEngine, modifies: Vec<Modify>) -> EngineResult<()> {
    // Acquire every CF in a stable order before applying anything. Readers can
    // therefore observe either the old batch or the new batch, never a partial
    // cross-CF write.
    let mut trees = engine
        .cf_contents
        .iter()
        .map(|tree| tree.write().unwrap())
        .collect::<Vec<_>>();

    let cf_index = |cf: CfName| {
        engine
            .cf_names
            .iter()
            .position(|name| *name == cf)
            .unwrap_or_else(|| panic!("CF {} does not exist", cf))
    };

    for modify in modifies {
        match modify {
            Modify::Delete(cf, k) => {
                trees[cf_index(cf)].remove(&k);
            }
            Modify::Put(cf, k, mut v) => {
                // The production commit log consumes the lock while applying
                // the commit. Preserve its inline value in this generic test
                // representation before the following lock deletion removes
                // the only other copy.
                if cf == CF_WRITE {
                    if let Ok(write) = WriteRef::parse(&v) {
                        if write.write_type == WriteType::Put && write.short_value.is_none() {
                            let lock_key = k.clone().truncate_ts().unwrap();
                            if let Some(lock_value) = trees[cf_index(CF_LOCK)].get(&lock_key) {
                                if let tikv_util::Either::Left(lock) =
                                    txn_types::parse_lock(lock_value).unwrap()
                                {
                                    if let Some(short_value) = lock.short_value {
                                        let mut write = write.to_owned();
                                        write.short_value = Some(short_value);
                                        v = write.as_ref().to_bytes();
                                    }
                                }
                            }
                        }
                    }
                }
                trees[cf_index(cf)].insert(k, v);
            }
            Modify::PessimisticLock(k, lock) => {
                let v = lock.into_lock().to_bytes();
                trees[cf_index(CF_LOCK)].insert(k, v);
            }
            Modify::DeleteRange(cf, start_key, end_key, notify_only) => {
                if notify_only {
                    continue;
                }
                let tree = &mut trees[cf_index(cf)];
                let keys = tree
                    .range((Included(start_key), Excluded(end_key)))
                    .map(|(key, _)| key.clone())
                    .collect::<Vec<_>>();
                for key in keys {
                    tree.remove(&key);
                }
            }
        };
    }
    Ok(())
}

#[cfg(test)]
pub mod tests {
    use engine_traits::{ALL_CFS, IterOptions};
    use txn_types::Write;

    use super::{
        super::{CfStatistics, TEST_ENGINE_CFS, tests::*},
        *,
    };
    use crate::{Cursor, ScanMode};

    #[test]
    fn test_btree_engine() {
        let mut engine = BTreeEngine::new(TEST_ENGINE_CFS);
        test_base_curd_options(&mut engine)
    }

    #[test]
    fn test_linear_of_btree_engine() {
        let mut engine = BTreeEngine::default();
        test_linear(&mut engine);
    }

    #[test]
    fn test_statistic_of_btree_engine() {
        let mut engine = BTreeEngine::default();
        test_cfs_statistics(&mut engine);
    }

    #[test]
    fn test_bounds_of_btree_engine() {
        let mut engine = BTreeEngine::default();
        let test_data = vec![
            (b"a1".to_vec(), b"v1".to_vec()),
            (b"a3".to_vec(), b"v3".to_vec()),
            (b"a5".to_vec(), b"v5".to_vec()),
            (b"a7".to_vec(), b"v7".to_vec()),
        ];
        for (k, v) in &test_data {
            must_put(&engine, k.as_slice(), v.as_slice());
        }
        let snap = engine.snapshot(Default::default()).unwrap();
        let mut statistics = CfStatistics::default();

        // lower bound > upper bound, seek() returns false.
        let mut iter_op = IterOptions::default();
        iter_op.set_lower_bound(b"a7", 0);
        iter_op.set_upper_bound(b"a3", 0);
        let mut cursor = Cursor::new(
            snap.iter(CF_DEFAULT, iter_op).unwrap(),
            ScanMode::Forward,
            false,
        );
        assert!(!cursor.seek(&Key::from_raw(b"a5"), &mut statistics).unwrap());

        let mut iter_op = IterOptions::default();
        iter_op.set_lower_bound(b"a3", 0);
        iter_op.set_upper_bound(b"a7", 0);
        let mut cursor = Cursor::new(
            snap.iter(CF_DEFAULT, iter_op).unwrap(),
            ScanMode::Forward,
            false,
        );

        assert!(cursor.seek(&Key::from_raw(b"a5"), &mut statistics).unwrap());
        assert!(!cursor.seek(&Key::from_raw(b"a8"), &mut statistics).unwrap());
        assert!(!cursor.seek(&Key::from_raw(b"a0"), &mut statistics).unwrap());

        assert!(cursor.seek_to_last(&mut statistics));

        let mut ret = vec![];
        loop {
            ret.push((
                Key::from_encoded(cursor.key(&mut statistics).to_vec())
                    .to_raw()
                    .unwrap(),
                cursor.value(&mut statistics).to_vec(),
            ));

            if !cursor.prev(&mut statistics) {
                break;
            }
        }
        ret.sort();
        assert_eq!(ret, test_data[1..3].to_vec());
    }

    #[test]
    fn test_iterator() {
        let engine = BTreeEngine::default();
        let test_data = vec![
            (b"a1".to_vec(), b"v1".to_vec()),
            (b"a3".to_vec(), b"v3".to_vec()),
            (b"a5".to_vec(), b"v5".to_vec()),
            (b"a7".to_vec(), b"v7".to_vec()),
        ];
        for (k, v) in &test_data {
            must_put(&engine, k.as_slice(), v.as_slice());
        }

        let iter_op = IterOptions::default();
        let tree = engine.get_cf(CF_DEFAULT);
        let mut iter = BTreeEngineIterator::new(tree, iter_op);
        assert!(!iter.valid().unwrap());

        assert!(iter.seek_to_first().unwrap());
        assert_eq!(iter.key(), Key::from_raw(b"a1").as_encoded().as_slice());
        assert!(!iter.prev().unwrap());
        assert!(iter.next().unwrap());
        assert_eq!(iter.key(), Key::from_raw(b"a3").as_encoded().as_slice());

        assert!(iter.seek(&Key::from_raw(b"a1")).unwrap());
        assert_eq!(iter.key(), Key::from_raw(b"a1").as_encoded().as_slice());

        assert!(iter.seek_to_last().unwrap());
        assert_eq!(iter.key(), Key::from_raw(b"a7").as_encoded().as_slice());
        assert!(!iter.next().unwrap());
        assert!(iter.prev().unwrap());
        assert_eq!(iter.key(), Key::from_raw(b"a5").as_encoded().as_slice());

        assert!(iter.seek(&Key::from_raw(b"a7")).unwrap());
        assert_eq!(iter.key(), Key::from_raw(b"a7").as_encoded().as_slice());

        assert!(!iter.seek_for_prev(&Key::from_raw(b"a0")).unwrap());

        assert!(iter.seek_for_prev(&Key::from_raw(b"a1")).unwrap());
        assert_eq!(iter.key(), Key::from_raw(b"a1").as_encoded().as_slice());

        assert!(iter.seek_for_prev(&Key::from_raw(b"a8")).unwrap());
        assert_eq!(iter.key(), Key::from_raw(b"a7").as_encoded().as_slice());
    }

    #[test]
    fn test_get_not_exist_cf() {
        let engine = BTreeEngine::new(&[]);
        ::panic_hook::recover_safe(|| engine.get_cf("not_exist_cf")).unwrap_err();
    }

    #[test]
    fn test_projected_kvengine_uses_api_key_layout() {
        fn check(api_version: ApiVersion, raw_key: &[u8]) {
            let mut engine = BTreeEngine::new_with_api_version(ALL_CFS, api_version);
            let commit_ts = 2.into();
            let value = b"value".to_vec();
            write_modifies(
                &engine,
                vec![Modify::Put(
                    CF_WRITE,
                    Key::from_raw(raw_key).append_ts(commit_ts),
                    Write::new(WriteType::Put, 1.into(), Some(value.clone()))
                        .as_ref()
                        .to_bytes(),
                )],
            )
            .unwrap();

            let snapshot = engine.snapshot(Default::default()).unwrap();
            let item = snapshot.get_kvengine_snap().unwrap().core.get(
                kvengine::WRITE_CF,
                raw_key,
                commit_ts.into_inner(),
            );
            assert_eq!(item.get_value(), value);
        }

        check(ApiVersion::V1, b"key");
        check(ApiVersion::V1ttl, b"key");

        let mut api_v2_key = vec![b'x', 0, 0, 1];
        api_v2_key.extend_from_slice(b"key");
        check(ApiVersion::V2, &api_v2_key);
    }

    #[test]
    fn test_trigger_not_leader_rejects_snapshot() {
        let mut engine = BTreeEngine::default();
        engine.trigger_not_leader();

        let error = engine.snapshot(Default::default()).unwrap_err();
        assert!(matches!(
            *error.0,
            EngineErrorInner::Request(ref request) if request.has_not_leader()
        ));
    }
}