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 ff5e7414560478cd4f051eb3252c81aca880fa02..d777c453b5ce3b486c3845073c2095630da637bc 100644
--- a/services/ui/ws/window_server_test_impl.cc
+++ b/services/ui/ws/window_server_test_impl.cc
@@ -30,28 +30,43 @@ 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) && WindowHasValidFrame(window)) {
-    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.";

[fwang, 2017/02/24 08:51:15]

This is maybe not a big concern for now (the function is only called in two
places and in separate tests) but maybe we should:
1) execute callback with a 0 value (or the current count) here.
2) or more complex: store the callback in the map too, execute the old callback
with the current count here and start again with the new callback.
3) or even more complex: accumulate the callbacks in the map and execute all of
them when the root count is reached.

WDYT?

+    return;
+  }
+
+  painted_window_roots_[client_name] = 0;
   WindowTree* tree = window_server_->GetTreeWithClientName(client_name);
   if (tree) {
     for (const ServerWindow* window : tree->roots()) {
       if (WindowHasValidFrame(window)) {
-        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(