1use crate::{
2 http::{self, HttpClient, Request, Response},
3 Client, Connection, Credentials, EstablishConnectionError, UserStore,
4};
5use anyhow::{anyhow, Result};
6use futures::{future::BoxFuture, stream::BoxStream, Future, StreamExt};
7use gpui::{executor, ModelHandle, TestAppContext};
8use parking_lot::Mutex;
9use rpc::{proto, ConnectionId, Peer, Receipt, TypedEnvelope};
10use std::{fmt, rc::Rc, sync::Arc};
11
12pub struct FakeServer {
13 peer: Arc<Peer>,
14 state: Arc<Mutex<FakeServerState>>,
15 user_id: u64,
16 executor: Rc<executor::Foreground>,
17}
18
19#[derive(Default)]
20struct FakeServerState {
21 incoming: Option<BoxStream<'static, Box<dyn proto::AnyTypedEnvelope>>>,
22 connection_id: Option<ConnectionId>,
23 forbid_connections: bool,
24 auth_count: usize,
25 access_token: usize,
26}
27
28impl FakeServer {
29 pub async fn for_client(
30 client_user_id: u64,
31 client: &mut Arc<Client>,
32 cx: &TestAppContext,
33 ) -> Self {
34 let server = Self {
35 peer: Peer::new(),
36 state: Default::default(),
37 user_id: client_user_id,
38 executor: cx.foreground(),
39 };
40
41 Arc::get_mut(client)
42 .unwrap()
43 .override_authenticate({
44 let state = Arc::downgrade(&server.state);
45 move |cx| {
46 let state = state.clone();
47 cx.spawn(move |_| async move {
48 let state = state.upgrade().ok_or_else(|| anyhow!("server dropped"))?;
49 let mut state = state.lock();
50 state.auth_count += 1;
51 let access_token = state.access_token.to_string();
52 Ok(Credentials {
53 user_id: client_user_id,
54 access_token,
55 })
56 })
57 }
58 })
59 .override_establish_connection({
60 let peer = Arc::downgrade(&server.peer).clone();
61 let state = Arc::downgrade(&server.state);
62 move |credentials, cx| {
63 let peer = peer.clone();
64 let state = state.clone();
65 let credentials = credentials.clone();
66 cx.spawn(move |cx| async move {
67 let state = state.upgrade().ok_or_else(|| anyhow!("server dropped"))?;
68 let peer = peer.upgrade().ok_or_else(|| anyhow!("server dropped"))?;
69 if state.lock().forbid_connections {
70 Err(EstablishConnectionError::Other(anyhow!(
71 "server is forbidding connections"
72 )))?
73 }
74
75 assert_eq!(credentials.user_id, client_user_id);
76
77 if credentials.access_token != state.lock().access_token.to_string() {
78 Err(EstablishConnectionError::Unauthorized)?
79 }
80
81 let (client_conn, server_conn, _) = Connection::in_memory(cx.background());
82 let (connection_id, io, incoming) =
83 peer.add_test_connection(server_conn, cx.background()).await;
84 cx.background().spawn(io).detach();
85 let mut state = state.lock();
86 state.connection_id = Some(connection_id);
87 state.incoming = Some(incoming);
88 Ok(client_conn)
89 })
90 }
91 });
92
93 client
94 .authenticate_and_connect(false, &cx.to_async())
95 .await
96 .unwrap();
97 server
98 }
99
100 pub fn disconnect(&self) {
101 self.peer.disconnect(self.connection_id());
102 let mut state = self.state.lock();
103 state.connection_id.take();
104 state.incoming.take();
105 }
106
107 pub fn auth_count(&self) -> usize {
108 self.state.lock().auth_count
109 }
110
111 pub fn roll_access_token(&self) {
112 self.state.lock().access_token += 1;
113 }
114
115 pub fn forbid_connections(&self) {
116 self.state.lock().forbid_connections = true;
117 }
118
119 pub fn allow_connections(&self) {
120 self.state.lock().forbid_connections = false;
121 }
122
123 pub fn send<T: proto::EnvelopedMessage>(&self, message: T) {
124 self.peer.send(self.connection_id(), message).unwrap();
125 }
126
127 pub async fn receive<M: proto::EnvelopedMessage>(&self) -> Result<TypedEnvelope<M>> {
128 self.executor.start_waiting();
129 let message = self
130 .state
131 .lock()
132 .incoming
133 .as_mut()
134 .expect("not connected")
135 .next()
136 .await
137 .ok_or_else(|| anyhow!("other half hung up"))?;
138 self.executor.finish_waiting();
139 let type_name = message.payload_type_name();
140 Ok(*message
141 .into_any()
142 .downcast::<TypedEnvelope<M>>()
143 .unwrap_or_else(|_| {
144 panic!(
145 "fake server received unexpected message type: {:?}",
146 type_name
147 );
148 }))
149 }
150
151 pub async fn respond<T: proto::RequestMessage>(
152 &self,
153 receipt: Receipt<T>,
154 response: T::Response,
155 ) {
156 self.peer.respond(receipt, response).unwrap()
157 }
158
159 fn connection_id(&self) -> ConnectionId {
160 self.state.lock().connection_id.expect("not connected")
161 }
162
163 pub async fn build_user_store(
164 &self,
165 client: Arc<Client>,
166 cx: &mut TestAppContext,
167 ) -> ModelHandle<UserStore> {
168 let http_client = FakeHttpClient::with_404_response();
169 let user_store = cx.add_model(|cx| UserStore::new(client, http_client, cx));
170 assert_eq!(
171 self.receive::<proto::GetUsers>()
172 .await
173 .unwrap()
174 .payload
175 .user_ids,
176 &[self.user_id]
177 );
178 user_store
179 }
180}
181
182pub struct FakeHttpClient {
183 handler: Box<
184 dyn 'static
185 + Send
186 + Sync
187 + Fn(Request) -> BoxFuture<'static, Result<Response, http::Error>>,
188 >,
189}
190
191impl FakeHttpClient {
192 pub fn new<Fut, F>(handler: F) -> Arc<dyn HttpClient>
193 where
194 Fut: 'static + Send + Future<Output = Result<Response, http::Error>>,
195 F: 'static + Send + Sync + Fn(Request) -> Fut,
196 {
197 Arc::new(Self {
198 handler: Box::new(move |req| Box::pin(handler(req))),
199 })
200 }
201
202 pub fn with_404_response() -> Arc<dyn HttpClient> {
203 Self::new(|_| async move {
204 Ok(isahc::Response::builder()
205 .status(404)
206 .body(Default::default())
207 .unwrap())
208 })
209 }
210}
211
212impl fmt::Debug for FakeHttpClient {
213 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
214 f.debug_struct("FakeHttpClient").finish()
215 }
216}
217
218impl HttpClient for FakeHttpClient {
219 fn send<'a>(&'a self, req: Request) -> BoxFuture<'a, Result<Response, crate::http::Error>> {
220 let future = (self.handler)(req);
221 Box::pin(async move { future.await.map(Into::into) })
222 }
223}