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