Richard Feldman/auto (#16696)91ffa022 years agoHistory
992 lines · 34.4 KB · Rust
RawDownload
use crate::{    auth::split_dev_server_token,    db::{tests::TestDb, NewUserParams, UserId},    executor::Executor,    rpc::{Principal, Server, ZedVersion, CLEANUP_TIMEOUT, RECONNECT_TIMEOUT},    AppState, Config, RateLimiter,};use anyhow::anyhow;use call::ActiveCall;use channel::{ChannelBuffer, ChannelStore};use client::{    self, proto::PeerId, ChannelId, Client, Connection, Credentials, EstablishConnectionError,    UserStore,};use clock::FakeSystemClock;use collab_ui::channel_view::ChannelView;use collections::{HashMap, HashSet};use fs::FakeFs;use futures::{channel::oneshot, StreamExt as _};use git::GitHostingProviderRegistry;use gpui::{BackgroundExecutor, Context, Model, Task, TestAppContext, View, VisualTestContext};use http_client::FakeHttpClient;use language::LanguageRegistry;use node_runtime::FakeNodeRuntime;use notifications::NotificationStore;use parking_lot::Mutex;use project::{Project, WorktreeId};use remote::SshSession;use rpc::{    proto::{self, ChannelRole},    RECEIVE_TIMEOUT,};use semantic_version::SemanticVersion;use serde_json::json;use session::{AppSession, Session};use settings::SettingsStore;use std::{    cell::{Ref, RefCell, RefMut},    env,    ops::{Deref, DerefMut},    path::Path,    sync::{        atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},        Arc,    },};use workspace::{Workspace, WorkspaceStore}; pub struct TestServer {    pub app_state: Arc<AppState>,    pub test_live_kit_server: Arc<live_kit_client::TestServer>,    server: Arc<Server>,    next_github_user_id: i32,    connection_killers: Arc<Mutex<HashMap<PeerId, Arc<AtomicBool>>>>,    forbid_connections: Arc<AtomicBool>,    _test_db: TestDb,} pub struct TestClient {    pub username: String,    pub app_state: Arc<workspace::AppState>,    channel_store: Model<ChannelStore>,    notification_store: Model<NotificationStore>,    state: RefCell<TestClientState>,} #[derive(Default)]struct TestClientState {    local_projects: Vec<Model<Project>>,    dev_server_projects: Vec<Model<Project>>,    buffers: HashMap<Model<Project>, HashSet<Model<language::Buffer>>>,    channel_buffers: HashSet<Model<ChannelBuffer>>,} pub struct ContactsSummary {    pub current: Vec<String>,    pub outgoing_requests: Vec<String>,    pub incoming_requests: Vec<String>,} impl TestServer {    pub async fn start(deterministic: BackgroundExecutor) -> Self {        static NEXT_LIVE_KIT_SERVER_ID: AtomicUsize = AtomicUsize::new(0);         let use_postgres = env::var("USE_POSTGRES").ok();        let use_postgres = use_postgres.as_deref();        let test_db = if use_postgres == Some("true") || use_postgres == Some("1") {            TestDb::postgres(deterministic.clone())        } else {            TestDb::sqlite(deterministic.clone())        };        let live_kit_server_id = NEXT_LIVE_KIT_SERVER_ID.fetch_add(1, SeqCst);        let live_kit_server = live_kit_client::TestServer::create(            format!("http://livekit.{}.test", live_kit_server_id),            format!("devkey-{}", live_kit_server_id),            format!("secret-{}", live_kit_server_id),            deterministic.clone(),        )        .unwrap();        let executor = Executor::Deterministic(deterministic.clone());        let app_state = Self::build_app_state(&test_db, &live_kit_server, executor.clone()).await;        let epoch = app_state            .db            .create_server(&app_state.config.zed_environment)            .await            .unwrap();        let server = Server::new(epoch, app_state.clone());        server.start().await.unwrap();        // Advance clock to ensure the server's cleanup task is finished.        deterministic.advance_clock(CLEANUP_TIMEOUT);        Self {            app_state,            server,            connection_killers: Default::default(),            forbid_connections: Default::default(),            next_github_user_id: 0,            _test_db: test_db,            test_live_kit_server: live_kit_server,        }    }     pub async fn start2(        cx_a: &mut TestAppContext,        cx_b: &mut TestAppContext,    ) -> (TestServer, TestClient, TestClient, ChannelId) {        let mut server = Self::start(cx_a.executor()).await;        let client_a = server.create_client(cx_a, "user_a").await;        let client_b = server.create_client(cx_b, "user_b").await;        let channel_id = server            .make_channel(                "test-channel",                None,                (&client_a, cx_a),                &mut [(&client_b, cx_b)],            )            .await;        cx_a.run_until_parked();         (server, client_a, client_b, channel_id)    }     pub async fn start1(cx: &mut TestAppContext) -> (TestServer, TestClient) {        let mut server = Self::start(cx.executor().clone()).await;        let client = server.create_client(cx, "user_a").await;        (server, client)    }     pub async fn reset(&self) {        self.app_state.db.reset();        let epoch = self            .app_state            .db            .create_server(&self.app_state.config.zed_environment)            .await            .unwrap();        self.server.reset(epoch);    }     pub async fn create_client(&mut self, cx: &mut TestAppContext, name: &str) -> TestClient {        let fs = FakeFs::new(cx.executor());         cx.update(|cx| {            if cx.has_global::<SettingsStore>() {                panic!("Same cx used to create two test clients")            }            let settings = SettingsStore::test(cx);            cx.set_global(settings);            release_channel::init(SemanticVersion::default(), cx);            client::init_settings(cx);        });         let clock = Arc::new(FakeSystemClock::default());        let http = FakeHttpClient::with_404_response();        let user_id = if let Ok(Some(user)) = self.app_state.db.get_user_by_github_login(name).await        {            user.id        } else {            let github_user_id = self.next_github_user_id;            self.next_github_user_id += 1;            self.app_state                .db                .create_user(                    &format!("{name}@example.com"),                    false,                    NewUserParams {                        github_login: name.into(),                        github_user_id,                    },                )                .await                .expect("creating user failed")                .user_id        };        let client_name = name.to_string();        let mut client = cx.update(|cx| Client::new(clock, http.clone(), cx));        let server = self.server.clone();        let db = self.app_state.db.clone();        let connection_killers = self.connection_killers.clone();        let forbid_connections = self.forbid_connections.clone();         Arc::get_mut(&mut client)            .unwrap()            .set_id(user_id.to_proto())            .override_authenticate(move |cx| {                cx.spawn(|_| async move {                    let access_token = "the-token".to_string();                    Ok(Credentials::User {                        user_id: user_id.to_proto(),                        access_token,                    })                })            })            .override_establish_connection(move |credentials, cx| {                assert_eq!(                    credentials,                    &Credentials::User {                        user_id: user_id.0 as u64,                        access_token: "the-token".into()                    }                );                 let server = server.clone();                let db = db.clone();                let connection_killers = connection_killers.clone();                let forbid_connections = forbid_connections.clone();                let client_name = client_name.clone();                cx.spawn(move |cx| async move {                    if forbid_connections.load(SeqCst) {                        Err(EstablishConnectionError::other(anyhow!(                            "server is forbidding connections"                        )))                    } else {                        let (client_conn, server_conn, killed) =                            Connection::in_memory(cx.background_executor().clone());                        let (connection_id_tx, connection_id_rx) = oneshot::channel();                        let user = db                            .get_user_by_id(user_id)                            .await                            .expect("retrieving user failed")                            .unwrap();                        cx.background_executor()                            .spawn(server.handle_connection(                                server_conn,                                client_name,                                Principal::User(user),                                ZedVersion(SemanticVersion::new(1, 0, 0)),                                None,                                Some(connection_id_tx),                                Executor::Deterministic(cx.background_executor().clone()),                            ))                            .detach();                        let connection_id = connection_id_rx.await.map_err(|e| {                            EstablishConnectionError::Other(anyhow!(                                "{} (is server shutting down?)",                                e                            ))                        })?;                        connection_killers                            .lock()                            .insert(connection_id.into(), killed);                        Ok(client_conn)                    }                })            });         let git_hosting_provider_registry = cx.update(GitHostingProviderRegistry::default_global);        git_hosting_provider_registry            .register_hosting_provider(Arc::new(git_hosting_providers::Github));         let user_store = cx.new_model(|cx| UserStore::new(client.clone(), cx));        let workspace_store = cx.new_model(|cx| WorkspaceStore::new(client.clone(), cx));        let language_registry = Arc::new(LanguageRegistry::test(cx.executor()));        let session = cx.new_model(|cx| AppSession::new(Session::test(), cx));        let app_state = Arc::new(workspace::AppState {            client: client.clone(),            user_store: user_store.clone(),            workspace_store,            languages: language_registry,            fs: fs.clone(),            build_window_options: |_, _| Default::default(),            node_runtime: FakeNodeRuntime::new(),            session,        });         let os_keymap = "keymaps/default-macos.json";         cx.update(|cx| {            theme::init(theme::LoadThemes::JustBase, cx);            Project::init(&client, cx);            client::init(&client, cx);            language::init(cx);            editor::init(cx);            workspace::init(app_state.clone(), cx);            call::init(client.clone(), user_store.clone(), cx);            channel::init(&client, user_store.clone(), cx);            notifications::init(client.clone(), user_store, cx);            collab_ui::init(&app_state, cx);            file_finder::init(cx);            menu::init();            dev_server_projects::init(client.clone(), cx);            settings::KeymapFile::load_asset(os_keymap, cx).unwrap();            language_model::LanguageModelRegistry::test(cx);            assistant::context_store::init(&client.clone().into());        });         client            .authenticate_and_connect(false, &cx.to_async())            .await            .unwrap();         let client = TestClient {            app_state,            username: name.to_string(),            channel_store: cx.read(ChannelStore::global).clone(),            notification_store: cx.read(NotificationStore::global).clone(),            state: Default::default(),        };        client.wait_for_current_user(cx).await;        client    }     pub async fn create_dev_server(        &self,        access_token: String,        cx: &mut TestAppContext,    ) -> TestClient {        cx.update(|cx| {            if cx.has_global::<SettingsStore>() {                panic!("Same cx used to create two test clients")            }            let settings = SettingsStore::test(cx);            cx.set_global(settings);            release_channel::init(SemanticVersion::default(), cx);            client::init_settings(cx);        });        let (dev_server_id, _) = split_dev_server_token(&access_token).unwrap();         let clock = Arc::new(FakeSystemClock::default());        let http = FakeHttpClient::with_404_response();        let mut client = cx.update(|cx| Client::new(clock, http.clone(), cx));        let server = self.server.clone();        let db = self.app_state.db.clone();        let connection_killers = self.connection_killers.clone();        let forbid_connections = self.forbid_connections.clone();        Arc::get_mut(&mut client)            .unwrap()            .set_id(1)            .set_dev_server_token(client::DevServerToken(access_token.clone()))            .override_establish_connection(move |credentials, cx| {                assert_eq!(                    credentials,                    &Credentials::DevServer {                        token: client::DevServerToken(access_token.to_string())                    }                );                 let server = server.clone();                let db = db.clone();                let connection_killers = connection_killers.clone();                let forbid_connections = forbid_connections.clone();                cx.spawn(move |cx| async move {                    if forbid_connections.load(SeqCst) {                        Err(EstablishConnectionError::other(anyhow!(                            "server is forbidding connections"                        )))                    } else {                        let (client_conn, server_conn, killed) =                            Connection::in_memory(cx.background_executor().clone());                        let (connection_id_tx, connection_id_rx) = oneshot::channel();                        let dev_server = db                            .get_dev_server(dev_server_id)                            .await                            .expect("retrieving dev_server failed");                        cx.background_executor()                            .spawn(server.handle_connection(                                server_conn,                                "dev-server".to_string(),                                Principal::DevServer(dev_server),                                ZedVersion(SemanticVersion::new(1, 0, 0)),                                None,                                Some(connection_id_tx),                                Executor::Deterministic(cx.background_executor().clone()),                            ))                            .detach();                        let connection_id = connection_id_rx.await.map_err(|e| {                            EstablishConnectionError::Other(anyhow!(                                "{} (is server shutting down?)",                                e                            ))                        })?;                        connection_killers                            .lock()                            .insert(connection_id.into(), killed);                        Ok(client_conn)                    }                })            });         let fs = FakeFs::new(cx.executor());        let user_store = cx.new_model(|cx| UserStore::new(client.clone(), cx));        let workspace_store = cx.new_model(|cx| WorkspaceStore::new(client.clone(), cx));        let language_registry = Arc::new(LanguageRegistry::test(cx.executor()));        let session = cx.new_model(|cx| AppSession::new(Session::test(), cx));        let app_state = Arc::new(workspace::AppState {            client: client.clone(),            user_store: user_store.clone(),            workspace_store,            languages: language_registry,            fs: fs.clone(),            build_window_options: |_, _| Default::default(),            node_runtime: FakeNodeRuntime::new(),            session,        });         cx.update(|cx| {            theme::init(theme::LoadThemes::JustBase, cx);            Project::init(&client, cx);            client::init(&client, cx);            language::init(cx);            editor::init(cx);            workspace::init(app_state.clone(), cx);            call::init(client.clone(), user_store.clone(), cx);            channel::init(&client, user_store.clone(), cx);            notifications::init(client.clone(), user_store, cx);            collab_ui::init(&app_state, cx);            file_finder::init(cx);            menu::init();            headless::init(                client.clone(),                headless::AppState {                    languages: app_state.languages.clone(),                    user_store: app_state.user_store.clone(),                    fs: fs.clone(),                    node_runtime: app_state.node_runtime.clone(),                },                cx,            )        })        .await        .unwrap();         TestClient {            app_state,            username: "dev-server".to_string(),            channel_store: cx.read(ChannelStore::global).clone(),            notification_store: cx.read(NotificationStore::global).clone(),            state: Default::default(),        }    }     pub fn disconnect_client(&self, peer_id: PeerId) {        self.connection_killers            .lock()            .remove(&peer_id)            .unwrap()            .store(true, SeqCst);    }     pub fn simulate_long_connection_interruption(        &self,        peer_id: PeerId,        deterministic: BackgroundExecutor,    ) {        self.forbid_connections();        self.disconnect_client(peer_id);        deterministic.advance_clock(RECEIVE_TIMEOUT + RECONNECT_TIMEOUT);        self.allow_connections();        deterministic.advance_clock(RECEIVE_TIMEOUT + RECONNECT_TIMEOUT);        deterministic.run_until_parked();    }     pub fn forbid_connections(&self) {        self.forbid_connections.store(true, SeqCst);    }     pub fn allow_connections(&self) {        self.forbid_connections.store(false, SeqCst);    }     pub async fn make_contacts(&self, clients: &mut [(&TestClient, &mut TestAppContext)]) {        for ix in 1..clients.len() {            let (left, right) = clients.split_at_mut(ix);            let (client_a, cx_a) = left.last_mut().unwrap();            for (client_b, cx_b) in right {                client_a                    .app_state                    .user_store                    .update(*cx_a, |store, cx| {                        store.request_contact(client_b.user_id().unwrap(), cx)                    })                    .await                    .unwrap();                cx_a.executor().run_until_parked();                client_b                    .app_state                    .user_store                    .update(*cx_b, |store, cx| {                        store.respond_to_contact_request(client_a.user_id().unwrap(), true, cx)                    })                    .await

Lines 1–500 of 992.