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