AppFlowy/frontend/rust-lib/flowy-task/tests/task_test/script.rs

use anyhow::Error;
use flowy_task::{Task, TaskContent, TaskDispatcher, TaskHandler, TaskId, TaskResult, TaskRunner, TaskState};
use futures::stream::FuturesUnordered;
use futures::StreamExt;
use lib_infra::async_trait::async_trait;
use lib_infra::future::BoxResultFuture;
use lib_infra::ref_map::RefCountValue;
use rand::Rng;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::oneshot::Receiver;
use tokio::sync::RwLock;

pub enum SearchScript {
    AddTask {
        task: Task,
    },
    AddTasks {
        tasks: Vec<Task>,
    },
    #[allow(dead_code)]
    Wait {
        millisecond: u64,
    },
    CancelTask {
        task_id: TaskId,
    },
    UnregisterHandler {
        handler_id: String,
    },
    AssertTaskStatus {
        task_id: TaskId,
        expected_status: TaskState,
    },
    AssertExecuteOrder {
        execute_order: Vec<u32>,
        rets: Vec<Receiver<TaskResult>>,
    },
}

pub struct SearchTest {
    scheduler: Arc<RwLock<TaskDispatcher>>,
}

impl SearchTest {
    pub async fn new() -> Self {
        let duration = Duration::from_millis(1000);
        let mut scheduler = TaskDispatcher::new(duration);
        scheduler.register_handler(Arc::new(MockTextTaskHandler()));
        scheduler.register_handler(Arc::new(MockBlobTaskHandler()));
        scheduler.register_handler(Arc::new(MockTimeoutTaskHandler()));

        let scheduler = Arc::new(RwLock::new(scheduler));
        tokio::spawn(TaskRunner::run(scheduler.clone()));

        Self { scheduler }
    }

    pub async fn next_task_id(&self) -> TaskId {
        self.scheduler.read().await.next_task_id()
    }

    pub async fn run_scripts(&self, scripts: Vec<SearchScript>) {
        for script in scripts {
            self.run_script(script).await;
        }
    }

    pub async fn run_script(&self, script: SearchScript) {
        match script {
            SearchScript::AddTask { task } => {
                self.scheduler.write().await.add_task(task);
            }
            SearchScript::CancelTask { task_id } => {
                self.scheduler.write().await.cancel_task(task_id);
            }
            SearchScript::AddTasks { tasks } => {
                let mut scheduler = self.scheduler.write().await;
                for task in tasks {
                    scheduler.add_task(task);
                }
            }
            SearchScript::Wait { millisecond } => {
                tokio::time::sleep(Duration::from_millis(millisecond)).await;
            }
            SearchScript::UnregisterHandler { handler_id } => {
                self.scheduler.write().await.unregister_handler(handler_id).await;
            }
            SearchScript::AssertTaskStatus {
                task_id,
                expected_status,
            } => {
                let status = self.scheduler.read().await.read_task(&task_id).unwrap().state().clone();
                assert_eq!(status, expected_status);
            }
            SearchScript::AssertExecuteOrder { execute_order, rets } => {
                let mut futures = FuturesUnordered::new();
                for ret in rets {
                    futures.push(ret);
                }
                let mut orders = vec![];
                while let Some(Ok(result)) = futures.next().await {
                    orders.push(result.id);
                    assert!(result.state.is_done());
                }
                assert_eq!(execute_order, orders);
            }
        }
    }
}

pub struct MockTextTaskHandler();
#[async_trait]
impl RefCountValue for MockTextTaskHandler {
    async fn did_remove(&self) {}
}

impl TaskHandler for MockTextTaskHandler {
    fn handler_id(&self) -> &str {
        "1"
    }

    fn run(&self, content: TaskContent) -> BoxResultFuture<(), Error> {
        let mut rng = rand::thread_rng();
        let millisecond = rng.gen_range(1..50);
        Box::pin(async move {
            match content {
                TaskContent::Text(_s) => {
                    tokio::time::sleep(Duration::from_millis(millisecond)).await;
                }
                TaskContent::Blob(_) => panic!("Only support text"),
            }
            Ok(())
        })
    }
}

pub fn make_text_background_task(task_id: TaskId, s: &str) -> (Task, Receiver<TaskResult>) {
    let mut task = Task::background("1", task_id, TaskContent::Text(s.to_owned()));
    let recv = task.recv.take().unwrap();
    (task, recv)
}

pub fn make_text_user_interactive_task(task_id: TaskId, s: &str) -> (Task, Receiver<TaskResult>) {
    let mut task = Task::user_interactive("1", task_id, TaskContent::Text(s.to_owned()));
    let recv = task.recv.take().unwrap();
    (task, recv)
}

pub struct MockBlobTaskHandler();
#[async_trait]
impl RefCountValue for MockBlobTaskHandler {
    async fn did_remove(&self) {}
}

impl TaskHandler for MockBlobTaskHandler {
    fn handler_id(&self) -> &str {
        "2"
    }

    fn run(&self, content: TaskContent) -> BoxResultFuture<(), Error> {
        Box::pin(async move {
            match content {
                TaskContent::Text(_) => panic!("Only support blob"),
                TaskContent::Blob(bytes) => {
                    let _msg = String::from_utf8(bytes).unwrap();
                    tokio::time::sleep(Duration::from_millis(20)).await;
                }
            }
            Ok(())
        })
    }
}

pub struct MockTimeoutTaskHandler();

impl TaskHandler for MockTimeoutTaskHandler {
    fn handler_id(&self) -> &str {
        "3"
    }

    fn run(&self, content: TaskContent) -> BoxResultFuture<(), Error> {
        Box::pin(async move {
            match content {
                TaskContent::Text(_) => panic!("Only support blob"),
                TaskContent::Blob(_bytes) => {
                    tokio::time::sleep(Duration::from_millis(2000)).await;
                }
            }
            Ok(())
        })
    }
}

pub fn make_timeout_task(task_id: TaskId) -> (Task, Receiver<TaskResult>) {
    let mut task = Task::background("3", task_id, TaskContent::Blob(vec![]));
    let recv = task.recv.take().unwrap();
    (task, recv)
}