Make WebDriver Get() command wait on pages loading before returning.

This makes using WebDriver significantly less racy. Also
refactors the message structure a little

components/compositing/constellation.rs

@@ -29,6 +29,7 @@ use msg::constellation_msg::{IFrameSandboxState, MozBrowserEvent, NavigationDire
 use msg::constellation_msg::{Key, KeyState, KeyModifiers, LoadData};
 use msg::constellation_msg::{SubpageId, WindowSizeData};
 use msg::constellation_msg::{self, ConstellationChan, Failure};
+use msg::constellation_msg::{WebDriverCommandMsg};
 use net_traits::{self, ResourceTask};
 use net_traits::image_cache_task::ImageCacheTask;
 use net_traits::storage_task::{StorageTask, StorageTaskMsg};
@@ -49,7 +50,7 @@ use util::geometry::PagePx;
 use util::opts;
 use util::task::spawn_named;
 use clipboard::ClipboardContext;
-use webdriver_traits::WebDriverScriptCommand;
+use webdriver_traits;
 
 /// Maintains the pipelines and navigation context and grants permission to composite.
 ///
@@ -122,6 +123,9 @@ pub struct Constellation<LTF, STF> {
 
     /// Means of accessing the clipboard
     clipboard_ctx: ClipboardContext,
+
+    /// Bits of state used to interact with the webdriver implementation
+    webdriver: WebDriverData
 }
 
 /// Stores the navigation context for a single frame in the frame tree.
@@ -185,6 +189,18 @@ pub struct SendableFrameTree {
     pub children: Vec<SendableFrameTree>,
 }
 
+struct WebDriverData {
+    load_channel: Option<Sender<webdriver_traits::LoadComplete>>
+}
+
+impl WebDriverData {
+    pub fn new() -> WebDriverData {
+        WebDriverData {
+            load_channel: None
+        }
+    }
+}
+
 #[derive(Clone, Copy)]
 enum ExitPipelineMode {
     Normal,
@@ -233,6 +249,7 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
                 },
                 phantom: PhantomData,
                 clipboard_ctx: ClipboardContext::new().unwrap(),
+                webdriver: WebDriverData::new()
             };
             constellation.run();
         });
@@ -376,7 +393,7 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
             // script, and reflow messages have been sent.
             ConstellationMsg::LoadComplete => {
                 debug!("constellation got load complete message");
-                self.compositor_proxy.send(CompositorMsg::LoadComplete);
+                self.handle_load_complete_msg()
             }
             // Handle a forward or back request
             ConstellationMsg::Navigate(pipeline_info, direction) => {
@@ -426,14 +443,9 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
                 };
                 sender.send(result).unwrap();
             }
-            ConstellationMsg::CompositePng(reply) => {
+            ConstellationMsg::WebDriverCommand(command) => {
-                self.compositor_proxy.send(CompositorMsg::CreatePng(reply));
-            }
-            ConstellationMsg::WebDriverCommand(pipeline_id,
-                                               command) => {
                 debug!("constellation got webdriver command message");
-                self.handle_webdriver_command_msg(pipeline_id,
+                self.handle_webdriver_msg(command);
-                                                  command);
             }
             ConstellationMsg::ViewportConstrained(pipeline_id, constraints) => {
                 debug!("constellation got viewport-constrained event message");
@@ -649,6 +661,14 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
         }
     }
 
+    fn handle_load_complete_msg(&mut self) {
+        self.compositor_proxy.send(CompositorMsg::LoadComplete);
+        if let Some(ref reply_chan) = self.webdriver.load_channel {
+            reply_chan.send(webdriver_traits::LoadComplete).unwrap();
+        }
+        self.webdriver.load_channel = None;
+    }
+
     fn handle_navigate_msg(&mut self,
                            pipeline_info: Option<(PipelineId, SubpageId)>,
                            direction: constellation_msg::NavigationDirection) {
@@ -814,15 +834,24 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
         }
     }
 
-    fn handle_webdriver_command_msg(&mut self,
+    fn handle_webdriver_msg(&mut self, msg: WebDriverCommandMsg) {
-                                    pipeline_id: PipelineId,
-                                    msg: WebDriverScriptCommand) {
         // Find the script channel for the given parent pipeline,
         // and pass the event to that script task.
-        let pipeline = self.pipeline(pipeline_id);
+        match msg {
-        let control_msg = ConstellationControlMsg::WebDriverCommand(pipeline_id, msg);
+            WebDriverCommandMsg::LoadUrl(pipeline_id, load_data, reply) => {
-        let ScriptControlChan(ref script_channel) = pipeline.script_chan;
+                self.handle_load_url_msg(pipeline_id, load_data);
-        script_channel.send(control_msg).unwrap();
+                self.webdriver.load_channel = Some(reply);
+            },
+            WebDriverCommandMsg::ScriptCommand(pipeline_id, cmd) => {
+                let pipeline = self.pipeline(pipeline_id);
+                let control_msg = ConstellationControlMsg::WebDriverScriptCommand(pipeline_id, cmd);
+                let ScriptControlChan(ref script_channel) = pipeline.script_chan;
+                script_channel.send(control_msg).unwrap();
+            },
+            WebDriverCommandMsg::TakeScreenshot(reply) => {
+                self.compositor_proxy.send(CompositorMsg::CreatePng(reply));
+            },
+        }
     }
 
     fn add_or_replace_pipeline_in_frame_tree(&mut self, frame_change: FrameChange) {

components/msg/constellation_msg.rs

@@ -18,7 +18,7 @@ use util::geometry::{PagePx, ViewportPx};
 use std::collections::HashMap;
 use std::sync::mpsc::{channel, Sender, Receiver};
 use style::viewport::ViewportConstraints;
-use webdriver_traits::WebDriverScriptCommand;
+use webdriver_traits::{WebDriverScriptCommand, LoadComplete};
 use url::Url;
 
 #[derive(Clone)]
@@ -237,11 +237,9 @@ pub enum Msg {
     /// Requests that the constellation retrieve the current contents of the clipboard
     GetClipboardContents(Sender<String>),
     /// Dispatch a webdriver command
-    WebDriverCommand(PipelineId, WebDriverScriptCommand),
+    WebDriverCommand(WebDriverCommandMsg),
     /// Notifies the constellation that the viewport has been constrained in some manner
     ViewportConstrained(PipelineId, ViewportConstraints),
-    /// Create a PNG of the window contents
-    CompositePng(Sender<Option<png::Image>>),
     /// Query the constellation to see if the current compositor output is stable
     IsReadyToSaveImage(HashMap<PipelineId, Epoch>),
     /// Notification that this iframe should be removed.
@@ -319,6 +317,11 @@ impl MozBrowserEvent {
     }
 }
 
+pub enum WebDriverCommandMsg {
+    LoadUrl(PipelineId, LoadData, Sender<LoadComplete>),
+    ScriptCommand(PipelineId, WebDriverScriptCommand),
+    TakeScreenshot(Sender<Option<png::Image>>)
+}
 
 /// Similar to net::resource_task::LoadData
 /// can be passed to LoadUrl to load a page with GET/POST

components/script/script_task.rs

@@ -726,7 +726,7 @@ impl ScriptTask {
                 self.handle_update_subpage_id(containing_pipeline_id, old_subpage_id, new_subpage_id),
             ConstellationControlMsg::FocusIFrame(containing_pipeline_id, subpage_id) =>
                 self.handle_focus_iframe_msg(containing_pipeline_id, subpage_id),
-            ConstellationControlMsg::WebDriverCommand(pipeline_id, msg) =>
+            ConstellationControlMsg::WebDriverScriptCommand(pipeline_id, msg) =>
                 self.handle_webdriver_msg(pipeline_id, msg),
             ConstellationControlMsg::TickAllAnimations(pipeline_id) =>
                 self.handle_tick_all_animations(pipeline_id),
@@ -801,8 +801,8 @@ impl ScriptTask {
     fn handle_webdriver_msg(&self, pipeline_id: PipelineId, msg: WebDriverScriptCommand) {
         let page = self.root_page();
         match msg {
-            WebDriverScriptCommand::EvaluateJS(script, reply) =>
+            WebDriverScriptCommand::ExecuteScript(script, reply) =>
-                webdriver_handlers::handle_evaluate_js(&page, pipeline_id, script, reply),
+                webdriver_handlers::handle_execute_script(&page, pipeline_id, script, reply),
             WebDriverScriptCommand::FindElementCSS(selector, reply) =>
                 webdriver_handlers::handle_find_element_css(&page, pipeline_id, selector, reply),
             WebDriverScriptCommand::FindElementsCSS(selector, reply) =>

components/script/webdriver_handlers.rs

@@ -35,7 +35,7 @@ fn find_node_by_unique_id(page: &Rc<Page>, pipeline: PipelineId, node_id: String
     None
 }
 
-pub fn handle_evaluate_js(page: &Rc<Page>, pipeline: PipelineId, eval: String, reply: Sender<Result<EvaluateJSReply, ()>>) {
+pub fn handle_execute_script(page: &Rc<Page>, pipeline: PipelineId, eval: String, reply: Sender<Result<EvaluateJSReply, ()>>) {
     let page = get_page(&*page, pipeline);
     let window = page.window().root();
     let cx = window.r().get_cx();

components/script_traits/lib.rs

@@ -98,7 +98,7 @@ pub enum ConstellationControlMsg {
     /// Set an iframe to be focused. Used when an element in an iframe gains focus.
     FocusIFrame(PipelineId, SubpageId),
     // Passes a webdriver command to the script task for execution
-    WebDriverCommand(PipelineId, WebDriverScriptCommand),
+    WebDriverScriptCommand(PipelineId, WebDriverScriptCommand),
     /// Notifies script task that all animations are done
     TickAllAnimations(PipelineId),
     /// Notifies script that a stylesheet has finished loading.

components/webdriver_server/lib.rs

@@ -19,7 +19,7 @@ extern crate rustc_serialize;
 extern crate uuid;
 extern crate webdriver_traits;
 
-use msg::constellation_msg::{ConstellationChan, LoadData, PipelineId, NavigationDirection};
+use msg::constellation_msg::{ConstellationChan, LoadData, PipelineId, NavigationDirection, WebDriverCommandMsg};
 use msg::constellation_msg::Msg as ConstellationMsg;
 use std::sync::mpsc::channel;
 use webdriver_traits::WebDriverScriptCommand;
@@ -111,10 +111,16 @@ impl Handler {
 
         let pipeline_id = self.get_root_pipeline();
 
+        let (sender, reciever) = channel();
+
         let load_data = LoadData::new(url);
         let ConstellationChan(ref const_chan) = self.constellation_chan;
-        const_chan.send(ConstellationMsg::LoadUrl(pipeline_id, load_data)).unwrap();
+        let cmd_msg = WebDriverCommandMsg::LoadUrl(pipeline_id, load_data, sender);
-        //TODO: Now we ought to wait until we get a load event
+        const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
+
+        //Wait to get a load event
+        reciever.recv().unwrap();
+
         Ok(WebDriverResponse::Void)
     }
 
@@ -135,8 +141,9 @@ impl Handler {
 
         let (sender, reciever) = channel();
         let ConstellationChan(ref const_chan) = self.constellation_chan;
-        const_chan.send(ConstellationMsg::WebDriverCommand(pipeline_id,
+        let cmd_msg = WebDriverCommandMsg::ScriptCommand(pipeline_id,
-                                                           WebDriverScriptCommand::GetTitle(sender))).unwrap();
+                                                         WebDriverScriptCommand::GetTitle(sender));
+        const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
         let value = reciever.recv().unwrap();
         Ok(WebDriverResponse::Generic(ValueResponse::new(value.to_json())))
     }
@@ -166,7 +173,8 @@ impl Handler {
         let (sender, reciever) = channel();
         let ConstellationChan(ref const_chan) = self.constellation_chan;
         let cmd = WebDriverScriptCommand::FindElementCSS(parameters.value.clone(), sender);
-        const_chan.send(ConstellationMsg::WebDriverCommand(pipeline_id, cmd)).unwrap();
+        let cmd_msg = WebDriverCommandMsg::ScriptCommand(pipeline_id, cmd);
+        const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
         match reciever.recv().unwrap() {
             Ok(value) => {
                 Ok(WebDriverResponse::Generic(ValueResponse::new(value.map(|x| WebElement::new(x).to_json()).to_json())))
@@ -187,7 +195,8 @@ impl Handler {
         let (sender, reciever) = channel();
         let ConstellationChan(ref const_chan) = self.constellation_chan;
         let cmd = WebDriverScriptCommand::FindElementsCSS(parameters.value.clone(), sender);
-        const_chan.send(ConstellationMsg::WebDriverCommand(pipeline_id, cmd)).unwrap();
+        let cmd_msg = WebDriverCommandMsg::ScriptCommand(pipeline_id, cmd);
+        const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
         match reciever.recv().unwrap() {
             Ok(value) => {
                 let resp_value: Vec<Json> = value.into_iter().map(
@@ -205,7 +214,8 @@ impl Handler {
         let (sender, reciever) = channel();
         let ConstellationChan(ref const_chan) = self.constellation_chan;
         let cmd = WebDriverScriptCommand::GetElementText(element.id.clone(), sender);
-        const_chan.send(ConstellationMsg::WebDriverCommand(pipeline_id, cmd)).unwrap();
+        let cmd_msg = WebDriverCommandMsg::ScriptCommand(pipeline_id, cmd);
+        const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
         match reciever.recv().unwrap() {
             Ok(value) => Ok(WebDriverResponse::Generic(ValueResponse::new(value.to_json()))),
             Err(_) => Err(WebDriverError::new(ErrorStatus::StaleElementReference,
@@ -219,7 +229,8 @@ impl Handler {
         let (sender, reciever) = channel();
         let ConstellationChan(ref const_chan) = self.constellation_chan;
         let cmd = WebDriverScriptCommand::GetActiveElement(sender);
-        const_chan.send(ConstellationMsg::WebDriverCommand(pipeline_id, cmd)).unwrap();
+        let cmd_msg = WebDriverCommandMsg::ScriptCommand(pipeline_id, cmd);
+        const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
         let value = reciever.recv().unwrap().map(|x| WebElement::new(x).to_json());
         Ok(WebDriverResponse::Generic(ValueResponse::new(value.to_json())))
     }
@@ -230,7 +241,8 @@ impl Handler {
         let (sender, reciever) = channel();
         let ConstellationChan(ref const_chan) = self.constellation_chan;
         let cmd = WebDriverScriptCommand::GetElementTagName(element.id.clone(), sender);
-        const_chan.send(ConstellationMsg::WebDriverCommand(pipeline_id, cmd)).unwrap();
+        let cmd_msg = WebDriverCommandMsg::ScriptCommand(pipeline_id, cmd);
+        const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
         match reciever.recv().unwrap() {
             Ok(value) => Ok(WebDriverResponse::Generic(ValueResponse::new(value.to_json()))),
             Err(_) => Err(WebDriverError::new(ErrorStatus::StaleElementReference,
@@ -253,8 +265,9 @@ impl Handler {
 
         let (sender, reciever) = channel();
         let ConstellationChan(ref const_chan) = self.constellation_chan;
-        const_chan.send(ConstellationMsg::WebDriverCommand(pipeline_id,
+        let cmd = WebDriverScriptCommand::ExecuteScript(script, sender);
-                                                           WebDriverScriptCommand::EvaluateJS(script, sender))).unwrap();
+        let cmd_msg = WebDriverCommandMsg::ScriptCommand(pipeline_id, cmd);
+        const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
 
         match reciever.recv().unwrap() {
             Ok(value) => Ok(WebDriverResponse::Generic(ValueResponse::new(value.to_json()))),
@@ -273,7 +286,8 @@ impl Handler {
         for _ in 0..iterations {
             let (sender, reciever) = channel();
             let ConstellationChan(ref const_chan) = self.constellation_chan;
-            const_chan.send(ConstellationMsg::CompositePng(sender)).unwrap();
+            let cmd_msg = WebDriverCommandMsg::TakeScreenshot(sender);
+            const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
 
             if let Some(x) = reciever.recv().unwrap() {
                 img = Some(x);

components/webdriver_traits/lib.rs

@@ -11,7 +11,7 @@ use rustc_serialize::json::{Json, ToJson};
 use std::sync::mpsc::Sender;
 
 pub enum WebDriverScriptCommand {
-    EvaluateJS(String, Sender<Result<EvaluateJSReply, ()>>),
+    ExecuteScript(String, Sender<Result<EvaluateJSReply, ()>>),
     FindElementCSS(String, Sender<Result<Option<String>, ()>>),
     FindElementsCSS(String, Sender<Result<Vec<String>, ()>>),
     GetActiveElement(Sender<Option<String>>),
@@ -40,3 +40,5 @@ impl ToJson for EvaluateJSReply {
         }
     }
 }
+
+pub struct LoadComplete;