test_server.rs

  1use crate::{
  2    db::{tests::TestDb, NewUserParams, UserId},
  3    executor::Executor,
  4    rpc::{Server, CLEANUP_TIMEOUT},
  5    AppState,
  6};
  7use anyhow::anyhow;
  8use call::ActiveCall;
  9use channel::{ChannelBuffer, ChannelStore};
 10use client::{
 11    self, proto::PeerId, Client, Connection, Credentials, EstablishConnectionError, UserStore,
 12};
 13use collections::{HashMap, HashSet};
 14use fs::FakeFs;
 15use futures::{channel::oneshot, StreamExt as _};
 16use gpui::{executor::Deterministic, ModelHandle, Task, TestAppContext, WindowHandle};
 17use language::LanguageRegistry;
 18use parking_lot::Mutex;
 19use project::{Project, WorktreeId};
 20use settings::SettingsStore;
 21use std::{
 22    cell::{Ref, RefCell, RefMut},
 23    env,
 24    ops::{Deref, DerefMut},
 25    path::Path,
 26    sync::{
 27        atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
 28        Arc,
 29    },
 30};
 31use util::http::FakeHttpClient;
 32use workspace::Workspace;
 33
 34pub struct TestServer {
 35    pub app_state: Arc<AppState>,
 36    pub test_live_kit_server: Arc<live_kit_client::TestServer>,
 37    server: Arc<Server>,
 38    connection_killers: Arc<Mutex<HashMap<PeerId, Arc<AtomicBool>>>>,
 39    forbid_connections: Arc<AtomicBool>,
 40    _test_db: TestDb,
 41}
 42
 43pub struct TestClient {
 44    pub username: String,
 45    pub app_state: Arc<workspace::AppState>,
 46    state: RefCell<TestClientState>,
 47}
 48
 49#[derive(Default)]
 50struct TestClientState {
 51    local_projects: Vec<ModelHandle<Project>>,
 52    remote_projects: Vec<ModelHandle<Project>>,
 53    buffers: HashMap<ModelHandle<Project>, HashSet<ModelHandle<language::Buffer>>>,
 54    channel_buffers: HashSet<ModelHandle<ChannelBuffer>>,
 55}
 56
 57pub struct ContactsSummary {
 58    pub current: Vec<String>,
 59    pub outgoing_requests: Vec<String>,
 60    pub incoming_requests: Vec<String>,
 61}
 62
 63impl TestServer {
 64    pub async fn start(deterministic: &Arc<Deterministic>) -> Self {
 65        static NEXT_LIVE_KIT_SERVER_ID: AtomicUsize = AtomicUsize::new(0);
 66
 67        let use_postgres = env::var("USE_POSTGRES").ok();
 68        let use_postgres = use_postgres.as_deref();
 69        let test_db = if use_postgres == Some("true") || use_postgres == Some("1") {
 70            TestDb::postgres(deterministic.build_background())
 71        } else {
 72            TestDb::sqlite(deterministic.build_background())
 73        };
 74        let live_kit_server_id = NEXT_LIVE_KIT_SERVER_ID.fetch_add(1, SeqCst);
 75        let live_kit_server = live_kit_client::TestServer::create(
 76            format!("http://livekit.{}.test", live_kit_server_id),
 77            format!("devkey-{}", live_kit_server_id),
 78            format!("secret-{}", live_kit_server_id),
 79            deterministic.build_background(),
 80        )
 81        .unwrap();
 82        let app_state = Self::build_app_state(&test_db, &live_kit_server).await;
 83        let epoch = app_state
 84            .db
 85            .create_server(&app_state.config.zed_environment)
 86            .await
 87            .unwrap();
 88        let server = Server::new(
 89            epoch,
 90            app_state.clone(),
 91            Executor::Deterministic(deterministic.build_background()),
 92        );
 93        server.start().await.unwrap();
 94        // Advance clock to ensure the server's cleanup task is finished.
 95        deterministic.advance_clock(CLEANUP_TIMEOUT);
 96        Self {
 97            app_state,
 98            server,
 99            connection_killers: Default::default(),
100            forbid_connections: Default::default(),
101            _test_db: test_db,
102            test_live_kit_server: live_kit_server,
103        }
104    }
105
106    pub async fn reset(&self) {
107        self.app_state.db.reset();
108        let epoch = self
109            .app_state
110            .db
111            .create_server(&self.app_state.config.zed_environment)
112            .await
113            .unwrap();
114        self.server.reset(epoch);
115    }
116
117    pub async fn create_client(&mut self, cx: &mut TestAppContext, name: &str) -> TestClient {
118        cx.update(|cx| {
119            if cx.has_global::<SettingsStore>() {
120                panic!("Same cx used to create two test clients")
121            }
122            cx.set_global(SettingsStore::test(cx));
123        });
124
125        let http = FakeHttpClient::with_404_response();
126        let user_id = if let Ok(Some(user)) = self.app_state.db.get_user_by_github_login(name).await
127        {
128            user.id
129        } else {
130            self.app_state
131                .db
132                .create_user(
133                    &format!("{name}@example.com"),
134                    false,
135                    NewUserParams {
136                        github_login: name.into(),
137                        github_user_id: 0,
138                        invite_count: 0,
139                    },
140                )
141                .await
142                .expect("creating user failed")
143                .user_id
144        };
145        let client_name = name.to_string();
146        let mut client = cx.read(|cx| Client::new(http.clone(), cx));
147        let server = self.server.clone();
148        let db = self.app_state.db.clone();
149        let connection_killers = self.connection_killers.clone();
150        let forbid_connections = self.forbid_connections.clone();
151
152        Arc::get_mut(&mut client)
153            .unwrap()
154            .set_id(user_id.0 as usize)
155            .override_authenticate(move |cx| {
156                cx.spawn(|_| async move {
157                    let access_token = "the-token".to_string();
158                    Ok(Credentials {
159                        user_id: user_id.0 as u64,
160                        access_token,
161                    })
162                })
163            })
164            .override_establish_connection(move |credentials, cx| {
165                assert_eq!(credentials.user_id, user_id.0 as u64);
166                assert_eq!(credentials.access_token, "the-token");
167
168                let server = server.clone();
169                let db = db.clone();
170                let connection_killers = connection_killers.clone();
171                let forbid_connections = forbid_connections.clone();
172                let client_name = client_name.clone();
173                cx.spawn(move |cx| async move {
174                    if forbid_connections.load(SeqCst) {
175                        Err(EstablishConnectionError::other(anyhow!(
176                            "server is forbidding connections"
177                        )))
178                    } else {
179                        let (client_conn, server_conn, killed) =
180                            Connection::in_memory(cx.background());
181                        let (connection_id_tx, connection_id_rx) = oneshot::channel();
182                        let user = db
183                            .get_user_by_id(user_id)
184                            .await
185                            .expect("retrieving user failed")
186                            .unwrap();
187                        cx.background()
188                            .spawn(server.handle_connection(
189                                server_conn,
190                                client_name,
191                                user,
192                                Some(connection_id_tx),
193                                Executor::Deterministic(cx.background()),
194                            ))
195                            .detach();
196                        let connection_id = connection_id_rx.await.unwrap();
197                        connection_killers
198                            .lock()
199                            .insert(connection_id.into(), killed);
200                        Ok(client_conn)
201                    }
202                })
203            });
204
205        let fs = FakeFs::new(cx.background());
206        let user_store = cx.add_model(|cx| UserStore::new(client.clone(), http, cx));
207        let channel_store =
208            cx.add_model(|cx| ChannelStore::new(client.clone(), user_store.clone(), cx));
209        let mut language_registry = LanguageRegistry::test();
210        language_registry.set_executor(cx.background());
211        let app_state = Arc::new(workspace::AppState {
212            client: client.clone(),
213            user_store: user_store.clone(),
214            channel_store: channel_store.clone(),
215            languages: Arc::new(language_registry),
216            fs: fs.clone(),
217            build_window_options: |_, _, _| Default::default(),
218            initialize_workspace: |_, _, _, _| Task::ready(Ok(())),
219            background_actions: || &[],
220        });
221
222        cx.update(|cx| {
223            theme::init((), cx);
224            Project::init(&client, cx);
225            client::init(&client, cx);
226            language::init(cx);
227            editor::init_settings(cx);
228            workspace::init(app_state.clone(), cx);
229            audio::init((), cx);
230            call::init(client.clone(), user_store.clone(), cx);
231            channel::init(&client);
232        });
233
234        client
235            .authenticate_and_connect(false, &cx.to_async())
236            .await
237            .unwrap();
238
239        let client = TestClient {
240            app_state,
241            username: name.to_string(),
242            state: Default::default(),
243        };
244        client.wait_for_current_user(cx).await;
245        client
246    }
247
248    pub fn disconnect_client(&self, peer_id: PeerId) {
249        self.connection_killers
250            .lock()
251            .remove(&peer_id)
252            .unwrap()
253            .store(true, SeqCst);
254    }
255
256    pub fn forbid_connections(&self) {
257        self.forbid_connections.store(true, SeqCst);
258    }
259
260    pub fn allow_connections(&self) {
261        self.forbid_connections.store(false, SeqCst);
262    }
263
264    pub async fn make_contacts(&self, clients: &mut [(&TestClient, &mut TestAppContext)]) {
265        for ix in 1..clients.len() {
266            let (left, right) = clients.split_at_mut(ix);
267            let (client_a, cx_a) = left.last_mut().unwrap();
268            for (client_b, cx_b) in right {
269                client_a
270                    .app_state
271                    .user_store
272                    .update(*cx_a, |store, cx| {
273                        store.request_contact(client_b.user_id().unwrap(), cx)
274                    })
275                    .await
276                    .unwrap();
277                cx_a.foreground().run_until_parked();
278                client_b
279                    .app_state
280                    .user_store
281                    .update(*cx_b, |store, cx| {
282                        store.respond_to_contact_request(client_a.user_id().unwrap(), true, cx)
283                    })
284                    .await
285                    .unwrap();
286            }
287        }
288    }
289
290    pub async fn make_channel(
291        &self,
292        channel: &str,
293        parent: Option<u64>,
294        admin: (&TestClient, &mut TestAppContext),
295        members: &mut [(&TestClient, &mut TestAppContext)],
296    ) -> u64 {
297        let (admin_client, admin_cx) = admin;
298        let channel_id = admin_client
299            .app_state
300            .channel_store
301            .update(admin_cx, |channel_store, cx| {
302                channel_store.create_channel(channel, parent, cx)
303            })
304            .await
305            .unwrap();
306
307        for (member_client, member_cx) in members {
308            admin_client
309                .app_state
310                .channel_store
311                .update(admin_cx, |channel_store, cx| {
312                    channel_store.invite_member(
313                        channel_id,
314                        member_client.user_id().unwrap(),
315                        false,
316                        cx,
317                    )
318                })
319                .await
320                .unwrap();
321
322            admin_cx.foreground().run_until_parked();
323
324            member_client
325                .app_state
326                .channel_store
327                .update(*member_cx, |channels, _| {
328                    channels.respond_to_channel_invite(channel_id, true)
329                })
330                .await
331                .unwrap();
332        }
333
334        channel_id
335    }
336
337    pub async fn make_channel_tree(
338        &self,
339        channels: &[(&str, Option<&str>)],
340        creator: (&TestClient, &mut TestAppContext),
341    ) -> Vec<u64> {
342        let mut observed_channels = HashMap::default();
343        let mut result = Vec::new();
344        for (channel, parent) in channels {
345            let id;
346            if let Some(parent) = parent {
347                if let Some(parent_id) = observed_channels.get(parent) {
348                    id = self
349                        .make_channel(channel, Some(*parent_id), (creator.0, creator.1), &mut [])
350                        .await;
351                } else {
352                    panic!(
353                        "Edge {}->{} referenced before {} was created",
354                        parent, channel, parent
355                    )
356                }
357            } else {
358                id = self
359                    .make_channel(channel, None, (creator.0, creator.1), &mut [])
360                    .await;
361            }
362
363            observed_channels.insert(channel, id);
364            result.push(id);
365        }
366
367        result
368    }
369
370    pub async fn create_room(&self, clients: &mut [(&TestClient, &mut TestAppContext)]) {
371        self.make_contacts(clients).await;
372
373        let (left, right) = clients.split_at_mut(1);
374        let (_client_a, cx_a) = &mut left[0];
375        let active_call_a = cx_a.read(ActiveCall::global);
376
377        for (client_b, cx_b) in right {
378            let user_id_b = client_b.current_user_id(*cx_b).to_proto();
379            active_call_a
380                .update(*cx_a, |call, cx| call.invite(user_id_b, None, cx))
381                .await
382                .unwrap();
383
384            cx_b.foreground().run_until_parked();
385            let active_call_b = cx_b.read(ActiveCall::global);
386            active_call_b
387                .update(*cx_b, |call, cx| call.accept_incoming(cx))
388                .await
389                .unwrap();
390        }
391    }
392
393    pub async fn build_app_state(
394        test_db: &TestDb,
395        fake_server: &live_kit_client::TestServer,
396    ) -> Arc<AppState> {
397        Arc::new(AppState {
398            db: test_db.db().clone(),
399            live_kit_client: Some(Arc::new(fake_server.create_api_client())),
400            config: Default::default(),
401        })
402    }
403}
404
405impl Deref for TestServer {
406    type Target = Server;
407
408    fn deref(&self) -> &Self::Target {
409        &self.server
410    }
411}
412
413impl Drop for TestServer {
414    fn drop(&mut self) {
415        self.server.teardown();
416        self.test_live_kit_server.teardown().unwrap();
417    }
418}
419
420impl Deref for TestClient {
421    type Target = Arc<Client>;
422
423    fn deref(&self) -> &Self::Target {
424        &self.app_state.client
425    }
426}
427
428impl TestClient {
429    pub fn fs(&self) -> &FakeFs {
430        self.app_state.fs.as_fake()
431    }
432
433    pub fn channel_store(&self) -> &ModelHandle<ChannelStore> {
434        &self.app_state.channel_store
435    }
436
437    pub fn user_store(&self) -> &ModelHandle<UserStore> {
438        &self.app_state.user_store
439    }
440
441    pub fn language_registry(&self) -> &Arc<LanguageRegistry> {
442        &self.app_state.languages
443    }
444
445    pub fn client(&self) -> &Arc<Client> {
446        &self.app_state.client
447    }
448
449    pub fn current_user_id(&self, cx: &TestAppContext) -> UserId {
450        UserId::from_proto(
451            self.app_state
452                .user_store
453                .read_with(cx, |user_store, _| user_store.current_user().unwrap().id),
454        )
455    }
456
457    pub async fn wait_for_current_user(&self, cx: &TestAppContext) {
458        let mut authed_user = self
459            .app_state
460            .user_store
461            .read_with(cx, |user_store, _| user_store.watch_current_user());
462        while authed_user.next().await.unwrap().is_none() {}
463    }
464
465    pub async fn clear_contacts(&self, cx: &mut TestAppContext) {
466        self.app_state
467            .user_store
468            .update(cx, |store, _| store.clear_contacts())
469            .await;
470    }
471
472    pub fn local_projects<'a>(&'a self) -> impl Deref<Target = Vec<ModelHandle<Project>>> + 'a {
473        Ref::map(self.state.borrow(), |state| &state.local_projects)
474    }
475
476    pub fn remote_projects<'a>(&'a self) -> impl Deref<Target = Vec<ModelHandle<Project>>> + 'a {
477        Ref::map(self.state.borrow(), |state| &state.remote_projects)
478    }
479
480    pub fn local_projects_mut<'a>(
481        &'a self,
482    ) -> impl DerefMut<Target = Vec<ModelHandle<Project>>> + 'a {
483        RefMut::map(self.state.borrow_mut(), |state| &mut state.local_projects)
484    }
485
486    pub fn remote_projects_mut<'a>(
487        &'a self,
488    ) -> impl DerefMut<Target = Vec<ModelHandle<Project>>> + 'a {
489        RefMut::map(self.state.borrow_mut(), |state| &mut state.remote_projects)
490    }
491
492    pub fn buffers_for_project<'a>(
493        &'a self,
494        project: &ModelHandle<Project>,
495    ) -> impl DerefMut<Target = HashSet<ModelHandle<language::Buffer>>> + 'a {
496        RefMut::map(self.state.borrow_mut(), |state| {
497            state.buffers.entry(project.clone()).or_default()
498        })
499    }
500
501    pub fn buffers<'a>(
502        &'a self,
503    ) -> impl DerefMut<Target = HashMap<ModelHandle<Project>, HashSet<ModelHandle<language::Buffer>>>> + 'a
504    {
505        RefMut::map(self.state.borrow_mut(), |state| &mut state.buffers)
506    }
507
508    pub fn channel_buffers<'a>(
509        &'a self,
510    ) -> impl DerefMut<Target = HashSet<ModelHandle<ChannelBuffer>>> + 'a {
511        RefMut::map(self.state.borrow_mut(), |state| &mut state.channel_buffers)
512    }
513
514    pub fn summarize_contacts(&self, cx: &TestAppContext) -> ContactsSummary {
515        self.app_state
516            .user_store
517            .read_with(cx, |store, _| ContactsSummary {
518                current: store
519                    .contacts()
520                    .iter()
521                    .map(|contact| contact.user.github_login.clone())
522                    .collect(),
523                outgoing_requests: store
524                    .outgoing_contact_requests()
525                    .iter()
526                    .map(|user| user.github_login.clone())
527                    .collect(),
528                incoming_requests: store
529                    .incoming_contact_requests()
530                    .iter()
531                    .map(|user| user.github_login.clone())
532                    .collect(),
533            })
534    }
535
536    pub async fn build_local_project(
537        &self,
538        root_path: impl AsRef<Path>,
539        cx: &mut TestAppContext,
540    ) -> (ModelHandle<Project>, WorktreeId) {
541        let project = self.build_empty_local_project(cx);
542        let (worktree, _) = project
543            .update(cx, |p, cx| {
544                p.find_or_create_local_worktree(root_path, true, cx)
545            })
546            .await
547            .unwrap();
548        worktree
549            .read_with(cx, |tree, _| tree.as_local().unwrap().scan_complete())
550            .await;
551        (project, worktree.read_with(cx, |tree, _| tree.id()))
552    }
553
554    pub fn build_empty_local_project(&self, cx: &mut TestAppContext) -> ModelHandle<Project> {
555        cx.update(|cx| {
556            Project::local(
557                self.client().clone(),
558                self.app_state.user_store.clone(),
559                self.app_state.languages.clone(),
560                self.app_state.fs.clone(),
561                cx,
562            )
563        })
564    }
565
566    pub async fn build_remote_project(
567        &self,
568        host_project_id: u64,
569        guest_cx: &mut TestAppContext,
570    ) -> ModelHandle<Project> {
571        let active_call = guest_cx.read(ActiveCall::global);
572        let room = active_call.read_with(guest_cx, |call, _| call.room().unwrap().clone());
573        room.update(guest_cx, |room, cx| {
574            room.join_project(
575                host_project_id,
576                self.app_state.languages.clone(),
577                self.app_state.fs.clone(),
578                cx,
579            )
580        })
581        .await
582        .unwrap()
583    }
584
585    pub fn build_workspace(
586        &self,
587        project: &ModelHandle<Project>,
588        cx: &mut TestAppContext,
589    ) -> WindowHandle<Workspace> {
590        cx.add_window(|cx| Workspace::new(0, project.clone(), self.app_state.clone(), cx))
591    }
592
593    pub async fn add_admin_to_channel(
594        &self,
595        user: (&TestClient, &mut TestAppContext),
596        channel: u64,
597        cx_self: &mut TestAppContext,
598    ) {
599        let (other_client, other_cx) = user;
600
601        self.app_state
602            .channel_store
603            .update(cx_self, |channel_store, cx| {
604                channel_store.invite_member(channel, other_client.user_id().unwrap(), true, cx)
605            })
606            .await
607            .unwrap();
608
609        cx_self.foreground().run_until_parked();
610
611        other_client
612            .app_state
613            .channel_store
614            .update(other_cx, |channels, _| {
615                channels.respond_to_channel_invite(channel, true)
616            })
617            .await
618            .unwrap();
619    }
620}
621
622impl Drop for TestClient {
623    fn drop(&mut self) {
624        self.app_state.client.teardown();
625    }
626}