Make mus_demo_unittests work with multiple root windows

services/ui/ws/window_server_test_impl.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "services/ui/ws/window_server_test_impl.h"
 
 #include "services/ui/public/interfaces/window_tree.mojom.h"
 #include "services/ui/ws/server_window.h"
 #include "services/ui/ws/server_window_compositor_frame_sink_manager.h"
 #include "services/ui/ws/window_server.h"
@@ skipping 12 matching lines @@
 
 }  // namespace
 
 WindowServerTestImpl::WindowServerTestImpl(WindowServer* window_server)
     : window_server_(window_server) {}
 
 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(
       base::Bind(&WindowServerTestImpl::OnWindowPaint, base::Unretained(this),
                  client_name, std::move(callback)));
 }
 
 }  // namespace ws
 }  // namespace ui