harness_lsp/
stub_client.rs

1use async_trait::async_trait;
2use std::collections::HashMap;
3use std::sync::{Arc, Mutex};
4
5use crate::types::{
6    CancelSignal, LspClient, LspHoverResult, LspLocation, LspServerProfile, LspSymbolInfo,
7    Position1, ServerHandle, ServerState,
8};
9
10pub type HoverResponder = Arc<
11    dyn Fn(&str, Position1) -> Option<LspHoverResult> + Send + Sync,
12>;
13pub type LocationResponder = Arc<dyn Fn(&str, Position1) -> Vec<LspLocation> + Send + Sync>;
14pub type DocSymbolResponder = Arc<dyn Fn(&str) -> Vec<LspSymbolInfo> + Send + Sync>;
15pub type WorkspaceSymbolResponder = Arc<dyn Fn(&str) -> Vec<LspSymbolInfo> + Send + Sync>;
16
17#[derive(Clone, Default)]
18pub struct StubResponses {
19    pub hover: Option<HoverResponder>,
20    pub definition: Option<LocationResponder>,
21    pub references: Option<LocationResponder>,
22    pub document_symbol: Option<DocSymbolResponder>,
23    pub workspace_symbol: Option<WorkspaceSymbolResponder>,
24    pub implementation: Option<LocationResponder>,
25}
26
27#[derive(Clone, Default)]
28pub struct StubBehavior {
29    /// Force ensureServer to return a "starting" handle for the first N calls
30    /// for this language; ready on the (N+1)th.
31    pub starting_calls: u32,
32    /// Force responses by language.
33    pub responses: HashMap<String, StubResponses>,
34    /// Force any next call for the given op to throw.
35    pub throw_on: Option<(Option<String>, Option<String>, String)>,
36    /// Force the op to hang until the cancel signal fires.
37    pub hang_on: Option<String>,
38}
39
40pub struct StubLspClient {
41    behavior: StubBehavior,
42    call_counts: Mutex<HashMap<String, u32>>,
43    closed: Mutex<bool>,
44}
45
46impl StubLspClient {
47    pub fn new(behavior: StubBehavior) -> Self {
48        Self {
49            behavior,
50            call_counts: Mutex::new(HashMap::new()),
51            closed: Mutex::new(false),
52        }
53    }
54
55    pub fn is_closed(&self) -> bool {
56        *self.closed.lock().unwrap()
57    }
58
59    fn maybe_throw(&self, language: &str, op: &str) -> Result<(), String> {
60        if let Some((l_opt, op_opt, err_msg)) = &self.behavior.throw_on {
61            let l_match = l_opt.as_deref().map(|l| l == language).unwrap_or(true);
62            let op_match = op_opt.as_deref().map(|o| o == op).unwrap_or(true);
63            if l_match && op_match {
64                return Err(err_msg.clone());
65            }
66        }
67        Ok(())
68    }
69
70    async fn maybe_hang(&self, op: &str, cancel: &CancelSignal) -> Result<(), String> {
71        if let Some(hang_op) = &self.behavior.hang_on {
72            if hang_op == op {
73                let mut rx = cancel.clone();
74                loop {
75                    if *rx.borrow() {
76                        return Err("aborted".to_string());
77                    }
78                    if rx.changed().await.is_err() {
79                        return Err("aborted".to_string());
80                    }
81                    if *rx.borrow() {
82                        return Err("aborted".to_string());
83                    }
84                }
85            }
86        }
87        Ok(())
88    }
89}
90
91#[async_trait]
92impl LspClient for StubLspClient {
93    async fn ensure_server(
94        &self,
95        language: &str,
96        root: &str,
97        _profile: &LspServerProfile,
98    ) -> Result<ServerHandle, String> {
99        let key = format!("{}|{}", language, root);
100        let mut counts = self.call_counts.lock().unwrap();
101        let next = counts.get(&key).copied().unwrap_or(0) + 1;
102        counts.insert(key, next);
103        let state = if next <= self.behavior.starting_calls {
104            ServerState::Starting
105        } else {
106            ServerState::Ready
107        };
108        Ok(ServerHandle {
109            language: language.to_string(),
110            root: root.to_string(),
111            state,
112        })
113    }
114
115    async fn hover(
116        &self,
117        handle: &ServerHandle,
118        path: &str,
119        pos: Position1,
120        cancel: CancelSignal,
121    ) -> Result<Option<LspHoverResult>, String> {
122        self.maybe_throw(&handle.language, "hover")?;
123        self.maybe_hang("hover", &cancel).await?;
124        Ok(self
125            .behavior
126            .responses
127            .get(&handle.language)
128            .and_then(|r| r.hover.as_ref())
129            .and_then(|f| f(path, pos)))
130    }
131
132    async fn definition(
133        &self,
134        handle: &ServerHandle,
135        path: &str,
136        pos: Position1,
137        cancel: CancelSignal,
138    ) -> Result<Vec<LspLocation>, String> {
139        self.maybe_throw(&handle.language, "definition")?;
140        self.maybe_hang("definition", &cancel).await?;
141        Ok(self
142            .behavior
143            .responses
144            .get(&handle.language)
145            .and_then(|r| r.definition.as_ref())
146            .map(|f| f(path, pos))
147            .unwrap_or_default())
148    }
149
150    async fn references(
151        &self,
152        handle: &ServerHandle,
153        path: &str,
154        pos: Position1,
155        cancel: CancelSignal,
156    ) -> Result<Vec<LspLocation>, String> {
157        self.maybe_throw(&handle.language, "references")?;
158        self.maybe_hang("references", &cancel).await?;
159        Ok(self
160            .behavior
161            .responses
162            .get(&handle.language)
163            .and_then(|r| r.references.as_ref())
164            .map(|f| f(path, pos))
165            .unwrap_or_default())
166    }
167
168    async fn document_symbol(
169        &self,
170        handle: &ServerHandle,
171        path: &str,
172        cancel: CancelSignal,
173    ) -> Result<Vec<LspSymbolInfo>, String> {
174        self.maybe_throw(&handle.language, "documentSymbol")?;
175        self.maybe_hang("documentSymbol", &cancel).await?;
176        Ok(self
177            .behavior
178            .responses
179            .get(&handle.language)
180            .and_then(|r| r.document_symbol.as_ref())
181            .map(|f| f(path))
182            .unwrap_or_default())
183    }
184
185    async fn workspace_symbol(
186        &self,
187        handle: &ServerHandle,
188        query: &str,
189        cancel: CancelSignal,
190    ) -> Result<Vec<LspSymbolInfo>, String> {
191        self.maybe_throw(&handle.language, "workspaceSymbol")?;
192        self.maybe_hang("workspaceSymbol", &cancel).await?;
193        Ok(self
194            .behavior
195            .responses
196            .get(&handle.language)
197            .and_then(|r| r.workspace_symbol.as_ref())
198            .map(|f| f(query))
199            .unwrap_or_default())
200    }
201
202    async fn implementation(
203        &self,
204        handle: &ServerHandle,
205        path: &str,
206        pos: Position1,
207        cancel: CancelSignal,
208    ) -> Result<Vec<LspLocation>, String> {
209        self.maybe_throw(&handle.language, "implementation")?;
210        self.maybe_hang("implementation", &cancel).await?;
211        Ok(self
212            .behavior
213            .responses
214            .get(&handle.language)
215            .and_then(|r| r.implementation.as_ref())
216            .map(|f| f(path, pos))
217            .unwrap_or_default())
218    }
219
220    async fn close_session(&self) {
221        *self.closed.lock().unwrap() = true;
222    }
223}
harness_lsp/stub_client.rs

harness_lsp/
stub_client.rs
