Skip to main content

allwright/
client_tab_query.rs

1use crate::proto::tab_session_command::Command as TabCommand;
2use crate::proto::tab_session_event::Event as TabEvent;
3use crate::proto::{
4    CountElementsCommand, GetInnerTextCommand, GetTextContentCommand, HighlightElementsCommand,
5    TabSessionCommand, WaitForSelectorCommand,
6};
7
8use super::command::{
9    command_retry_options, count_result_from_event, highlight_result_from_event,
10};
11use super::selectors::normalize_selector_for_transport;
12use super::tab::ensure_tab_open;
13use super::types::{
14    CommandOptions, Error, HighlightOptions, HighlightResult, Result, Tab, TextResult,
15    WaitForSelectorOptions, WaitForSelectorResult,
16};
17
18impl Tab {
19    pub async fn count(&self, css_selector: impl Into<String>) -> Result<super::types::CountResult> {
20        self.count_with_options(css_selector, CommandOptions::default()).await
21    }
22
23    pub async fn count_with_options(
24        &self,
25        css_selector: impl Into<String>,
26        options: CommandOptions,
27    ) -> Result<super::types::CountResult> {
28        let css_selector = normalize_selector_for_transport(&css_selector.into());
29        let mut state = self.inner.state.lock().await;
30        let handle = self.ensure_handle(&mut state).await?;
31        ensure_tab_open(handle, &self.inner.session_id)?;
32
33        handle
34            .command_tx
35            .send(TabSessionCommand {
36                browser_session_id: self.inner.browser_session_id.clone(),
37                tab_session_id: self.inner.session_id.clone(),
38                command: Some(TabCommand::CountElements(CountElementsCommand {
39                    css_selector: css_selector.clone(),
40                    retry_options: command_retry_options(options.timeout_ms),
41                })),
42            })
43            .await
44            .map_err(|_| Error::new("failed to send CountElementsCommand"))?;
45
46        loop {
47            let event = handle
48                .events
49                .message()
50                .await?
51                .ok_or_else(|| Error::new("tab session closed while waiting for count result"))?;
52
53            match event.event {
54                Some(TabEvent::Attached(_)) => {}
55                Some(TabEvent::ElementCounted(counted)) => {
56                    return Ok(count_result_from_event(counted));
57                }
58                Some(TabEvent::Error(error)) => {
59                    return Err(Error::new(format!(
60                        "tab session error while counting locator {:?}: {}",
61                        css_selector,
62                        error.message,
63                    )));
64                }
65                Some(TabEvent::Closed(_)) => {
66                    handle.closed = true;
67                    return Err(Error::new(format!(
68                        "tab session {} closed while waiting for count result",
69                        self.inner.session_id
70                    )));
71                }
72                _ => {}
73            }
74        }
75    }
76
77    pub async fn highlight(&self, css_selector: impl Into<String>) -> Result<HighlightResult> {
78        self.highlight_with_options(css_selector, HighlightOptions::default())
79            .await
80    }
81
82    pub async fn highlight_with_options(
83        &self,
84        css_selector: impl Into<String>,
85        options: HighlightOptions,
86    ) -> Result<HighlightResult> {
87        let css_selector = normalize_selector_for_transport(&css_selector.into());
88        let mut state = self.inner.state.lock().await;
89        let handle = self.ensure_handle(&mut state).await?;
90        ensure_tab_open(handle, &self.inner.session_id)?;
91
92        handle
93            .command_tx
94            .send(TabSessionCommand {
95                browser_session_id: self.inner.browser_session_id.clone(),
96                tab_session_id: self.inner.session_id.clone(),
97                command: Some(TabCommand::HighlightElements(HighlightElementsCommand {
98                    css_selector: css_selector.clone(),
99                    duration_ms: options.duration_ms,
100                    retry_options: command_retry_options(options.timeout_ms),
101                })),
102            })
103            .await
104            .map_err(|_| Error::new("failed to send HighlightElementsCommand"))?;
105
106        loop {
107            let event = handle
108                .events
109                .message()
110                .await?
111                .ok_or_else(|| Error::new("tab session closed while waiting for highlight result"))?;
112
113            match event.event {
114                Some(TabEvent::Attached(_)) => {}
115                Some(TabEvent::ElementsHighlighted(highlighted)) => {
116                    return Ok(highlight_result_from_event(highlighted));
117                }
118                Some(TabEvent::Error(error)) => {
119                    return Err(Error::new(format!(
120                        "tab session error while highlighting locator {:?}: {}",
121                        css_selector,
122                        error.message,
123                    )));
124                }
125                Some(TabEvent::Closed(_)) => {
126                    handle.closed = true;
127                    return Err(Error::new(format!(
128                        "tab session {} closed while waiting for highlight result",
129                        self.inner.session_id
130                    )));
131                }
132                _ => {}
133            }
134        }
135    }
136
137    pub async fn text_content(&self, css_selector: impl Into<String>) -> Result<TextResult> {
138        self.text_content_with_options(css_selector, CommandOptions::default())
139            .await
140    }
141
142    pub async fn text_content_with_options(
143        &self,
144        css_selector: impl Into<String>,
145        options: CommandOptions,
146    ) -> Result<TextResult> {
147        self.read_text(normalize_selector_for_transport(&css_selector.into()), options, true)
148            .await
149    }
150
151    pub async fn inner_text(&self, css_selector: impl Into<String>) -> Result<TextResult> {
152        self.inner_text_with_options(css_selector, CommandOptions::default())
153            .await
154    }
155
156    pub async fn inner_text_with_options(
157        &self,
158        css_selector: impl Into<String>,
159        options: CommandOptions,
160    ) -> Result<TextResult> {
161        self.read_text(normalize_selector_for_transport(&css_selector.into()), options, false)
162            .await
163    }
164
165    pub async fn wait_for_selector(
166        &self,
167        css_selector: impl Into<String>,
168    ) -> Result<WaitForSelectorResult> {
169        self.wait_for_selector_with_options(css_selector, WaitForSelectorOptions::default())
170            .await
171    }
172
173    pub async fn wait_for_selector_with_options(
174        &self,
175        css_selector: impl Into<String>,
176        options: WaitForSelectorOptions,
177    ) -> Result<WaitForSelectorResult> {
178        let css_selector = normalize_selector_for_transport(&css_selector.into());
179        let mut state = self.inner.state.lock().await;
180        let handle = self.ensure_handle(&mut state).await?;
181        ensure_tab_open(handle, &self.inner.session_id)?;
182        handle
183            .command_tx
184            .send(TabSessionCommand {
185                browser_session_id: self.inner.browser_session_id.clone(),
186                tab_session_id: self.inner.session_id.clone(),
187                command: Some(TabCommand::WaitForSelector(WaitForSelectorCommand {
188                    css_selector: css_selector.clone(),
189                    visible: options.visible,
190                    retry_options: command_retry_options(options.timeout_ms),
191                })),
192            })
193            .await
194            .map_err(|_| Error::new("failed to send WaitForSelectorCommand"))?;
195        loop {
196            let event = handle
197                .events
198                .message()
199                .await?
200                .ok_or_else(|| Error::new("tab session closed while waiting for selector"))?;
201            match event.event {
202                Some(TabEvent::Attached(_)) => {}
203                Some(TabEvent::SelectorWaitSatisfied(waited)) => {
204                    return Ok(WaitForSelectorResult {
205                        selector: waited.css_selector,
206                        visible: waited.visible,
207                        note: waited.note,
208                    });
209                }
210                Some(TabEvent::Error(error)) => {
211                    return Err(Error::new(format!(
212                        "tab session error while waiting for locator {:?}: {}",
213                        css_selector,
214                        error.message,
215                    )));
216                }
217                Some(TabEvent::Closed(_)) => {
218                    handle.closed = true;
219                    return Err(Error::new(format!(
220                        "tab session {} closed while waiting for selector result",
221                        self.inner.session_id
222                    )));
223                }
224                _ => {}
225            }
226        }
227    }
228
229    async fn read_text(
230        &self,
231        css_selector: String,
232        options: CommandOptions,
233        text_content: bool,
234    ) -> Result<TextResult> {
235        let mut state = self.inner.state.lock().await;
236        let handle = self.ensure_handle(&mut state).await?;
237        ensure_tab_open(handle, &self.inner.session_id)?;
238        let command = if text_content {
239            TabCommand::GetTextContent(GetTextContentCommand {
240                css_selector,
241                retry_options: command_retry_options(options.timeout_ms),
242            })
243        } else {
244            TabCommand::GetInnerText(GetInnerTextCommand {
245                css_selector,
246                retry_options: command_retry_options(options.timeout_ms),
247            })
248        };
249        handle
250            .command_tx
251            .send(TabSessionCommand {
252                browser_session_id: self.inner.browser_session_id.clone(),
253                tab_session_id: self.inner.session_id.clone(),
254                command: Some(command),
255            })
256            .await
257            .map_err(|_| Error::new("failed to send text command"))?;
258        loop {
259            let event = handle
260                .events
261                .message()
262                .await?
263                .ok_or_else(|| Error::new("tab session closed while waiting for text result"))?;
264            match event.event {
265                Some(TabEvent::Attached(_)) => {}
266                Some(TabEvent::TextContentResolved(text)) => {
267                    return Ok(TextResult {
268                        selector: text.css_selector,
269                        text: text.text,
270                        note: text.note,
271                    });
272                }
273                Some(TabEvent::InnerTextResolved(text)) => {
274                    return Ok(TextResult {
275                        selector: text.css_selector,
276                        text: text.text,
277                        note: text.note,
278                    });
279                }
280                Some(TabEvent::Error(error)) => {
281                    return Err(Error::new(format!(
282                        "tab session error while reading text: {}",
283                        error.message
284                    )));
285                }
286                Some(TabEvent::Closed(_)) => {
287                    handle.closed = true;
288                    return Err(Error::new(format!(
289                        "tab session {} closed while waiting for text result",
290                        self.inner.session_id
291                    )));
292                }
293                _ => {}
294            }
295        }
296    }
297}