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