Make mus_demo_unittests work with multiple root windows

services/ui/ws/window_server_test_impl.cc

Index: services/ui/ws/window_server_test_impl.cc
diff --git a/services/ui/ws/window_server_test_impl.cc b/services/ui/ws/window_server_test_impl.cc
index 20bc295c323a9dc0c30776295bcc7b0f25cad68d..007ae339a476e3d7cb16a68a529cb5b673a1fb24 100644
--- a/services/ui/ws/window_server_test_impl.cc
+++ b/services/ui/ws/window_server_test_impl.cc
@@ -20,28 +20,44 @@ WindowServerTestImpl::~WindowServerTestImpl() {}
 
 void WindowServerTestImpl::OnWindowPaint(
     const std::string& name,
-    const EnsureClientHasDrawnWindowCallback& cb,
+    const EnsureClientHasDrawnRootWindowsCallback& cb,
     ServerWindow* window) {
   WindowTree* tree = window_server_->GetTreeWithClientName(name);
   if (!tree)
     return;
   if (tree->HasRoot(window) && window->compositor_frame_sink_manager()) {
-    cb.Run(true);
-    window_server_->SetPaintCallback(base::Callback<void(ServerWindow*)>());
+    painted_window_roots_[name]++;
+    if (painted_window_roots_[name] == tree->roots().size()) {
+      painted_window_roots_.erase(name);
+      cb.Run(tree->roots().size());
+      window_server_->SetPaintCallback(base::Callback<void(ServerWindow*)>());
+    }
   }
 }
 
-void WindowServerTestImpl::EnsureClientHasDrawnWindow(
+void WindowServerTestImpl::EnsureClientHasDrawnRootWindows(
     const std::string& client_name,
-    const EnsureClientHasDrawnWindowCallback& callback) {
+    const EnsureClientHasDrawnRootWindowsCallback& callback) {
+  if (painted_window_roots_.find(client_name) != painted_window_roots_.end()) {
+    LOG(ERROR) << "EnsureClientHasDrawnRootWindows is already being executed "
+                  "for that client name.";
+    callback.Run(0);

[fwang, 2017/03/01 17:24:34]

I added this for completeness, although it's actually not needed.

+    return;
+  }
+
+  painted_window_roots_[client_name] = 0;
   WindowTree* tree = window_server_->GetTreeWithClientName(client_name);
   if (tree) {
     for (const ServerWindow* window : tree->roots()) {
       if (window->compositor_frame_sink_manager()) {
-        callback.Run(true);
-        return;
+        painted_window_roots_[client_name]++;
       }
     }
+    if (painted_window_roots_[client_name] == tree->roots().size()) {
+      painted_window_roots_.erase(client_name);
+      callback.Run(tree->roots().size());
+      return;
+    }
   }
 
   window_server_->SetPaintCallback(