Improve the stability of tests run using WebDriver.

* Wait for the correct pipeline to return a LoadComplete message
  before signalling a load is complete, and ensure that the root
  pipeline is the one corresponding to the active document of the top
  level browsing context, even if this has not yet painted.

* Ensure that TakeScreenshot operates on the correct pipeline

* Reset the screenshot ready flag whenever we decide that we are ready
  to take a screenshot.

components/compositing/compositor.rs

@@ -1281,6 +1281,9 @@ impl<Window: WindowMethods> IOCompositor<Window> {
                 // Constellation has replied at some point in the past
                 // that the current output image is stable and ready
                 // for saving.
+                // Reset the flag so that we check again in the future
+                // TODO: only reset this if we load a new document?
+                self.ready_to_save_state = ReadyState::Unknown;
                 true
             }
         }

components/compositing/constellation.rs

@@ -190,7 +190,7 @@ pub struct SendableFrameTree {
 }
 
 struct WebDriverData {
-    load_channel: Option<Sender<webdriver_msg::LoadComplete>>
+    load_channel: Option<(PipelineId, Sender<webdriver_msg::LoadComplete>)>
 }
 
 impl WebDriverData {
@@ -629,6 +629,10 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
     }
 
     fn handle_load_url_msg(&mut self, source_id: PipelineId, load_data: LoadData) {
+        self.load_url(source_id, load_data);
+    }
+
+    fn load_url(&mut self, source_id: PipelineId, load_data: LoadData) -> Option<PipelineId> {
         // If this load targets an iframe, its framing element may exist
         // in a separate script task than the framed document that initiated
         // the new load. The framing element must be notified about the
@@ -643,13 +647,14 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
                 script_channel.send(ConstellationControlMsg::Navigate(parent_pipeline_id,
                                                                       subpage_id,
                                                                       load_data)).unwrap();
+                Some(source_id)
             }
             None => {
                 // Make sure no pending page would be overridden.
                 for frame_change in &self.pending_frames {
                     if frame_change.old_pipeline_id == Some(source_id) {
                         // id that sent load msg is being changed already; abort
-                        return;
+                        return None;
                     }
                 }
 
@@ -665,6 +670,7 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
                 // Send message to ScriptTask that will suspend all timers
                 let old_pipeline = self.pipelines.get(&source_id).unwrap();
                 old_pipeline.freeze();
+                Some(new_pipeline_id)
             }
         }
     }
@@ -694,10 +700,17 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
         let forward = !self.mut_frame(frame_id).next.is_empty();
         let back = !self.mut_frame(frame_id).prev.is_empty();
         self.compositor_proxy.send(CompositorMsg::LoadComplete(back, forward));
-        if let Some(ref reply_chan) = self.webdriver.load_channel {
-            reply_chan.send(webdriver_msg::LoadComplete).unwrap();
+
+        let mut webdriver_reset = false;
+        if let Some((ref expected_pipeline_id, ref reply_chan)) = self.webdriver.load_channel {
+            if expected_pipeline_id == pipeline_id {
+                reply_chan.send(webdriver_msg::LoadComplete).unwrap();
+                webdriver_reset = true;
+            }
+        }
+        if webdriver_reset {
+            self.webdriver.load_channel = None;
         }
-        self.webdriver.load_channel = None;
     }
 
     fn handle_navigate_msg(&mut self,
@@ -821,10 +834,13 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
 
     fn handle_get_pipeline(&mut self, frame_id: Option<FrameId>,
                            resp_chan: Sender<Option<PipelineId>>) {
-        let pipeline_id = frame_id.or(self.root_frame_id).map(|frame_id| {
+        let current_pipeline_id = frame_id.or(self.root_frame_id).map(|frame_id| {
             let frame = self.frames.get(&frame_id).unwrap();
             frame.current
         });
+        let pipeline_id = self.pending_frames.iter().rev()
+            .find(|x| x.old_pipeline_id == current_pipeline_id)
+            .map(|x| x.new_pipeline_id).or(current_pipeline_id);
         resp_chan.send(pipeline_id).unwrap();
     }
 
@@ -880,17 +896,27 @@ impl<LTF: LayoutTaskFactory, STF: ScriptTaskFactory> Constellation<LTF, STF> {
         // and pass the event to that script task.
         match msg {
             WebDriverCommandMsg::LoadUrl(pipeline_id, load_data, reply) => {
-                self.handle_load_url_msg(pipeline_id, load_data);
-                self.webdriver.load_channel = Some(reply);
+                let new_pipeline_id = self.load_url(pipeline_id, load_data);
+                if let Some(id) = new_pipeline_id {
+                    self.webdriver.load_channel = Some((id, 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));
+            WebDriverCommandMsg::TakeScreenshot(pipeline_id, reply) => {
+                let current_pipeline_id = self.root_frame_id.map(|frame_id| {
+                    let frame = self.frames.get(&frame_id).unwrap();
+                    frame.current
+                });
+                if Some(pipeline_id) == current_pipeline_id {
+                    self.compositor_proxy.send(CompositorMsg::CreatePng(reply));
+                } else {
+                    reply.send(None).unwrap();
+                }
             },
         }
     }

components/msg/constellation_msg.rs

@@ -326,7 +326,7 @@ impl MozBrowserEvent {
 pub enum WebDriverCommandMsg {
     LoadUrl(PipelineId, LoadData, Sender<LoadComplete>),
     ScriptCommand(PipelineId, WebDriverScriptCommand),
-    TakeScreenshot(Sender<Option<png::Image>>)
+    TakeScreenshot(PipelineId, Sender<Option<png::Image>>)
 }
 
 /// Similar to net::resource_task::LoadData

components/webdriver_server/lib.rs

@@ -420,14 +420,15 @@ impl Handler {
 
     fn handle_take_screenshot(&self) -> WebDriverResult<WebDriverResponse> {
         let mut img = None;
+        let pipeline_id = try!(self.get_root_pipeline());
 
         let interval = 20;
         let iterations = 30_000 / interval;
 
         for _ in 0..iterations {
             let (sender, reciever) = channel();
             let ConstellationChan(ref const_chan) = self.constellation_chan;
-            let cmd_msg = WebDriverCommandMsg::TakeScreenshot(sender);
+            let cmd_msg = WebDriverCommandMsg::TakeScreenshot(pipeline_id, sender);
             const_chan.send(ConstellationMsg::WebDriverCommand(cmd_msg)).unwrap();
 
             if let Some(x) = reciever.recv().unwrap() {