MockColorChooser

components/test_runner/mock_color_chooser.cc

Index: components/test_runner/mock_color_chooser.cc
diff --git a/components/test_runner/mock_color_chooser.cc b/components/test_runner/mock_color_chooser.cc
index 48f426dd82e3786d81a15fdef4fc15aeea16196c..efa4d1d8973c0eff8f228bd80f9ef520d51d2607 100644
--- a/components/test_runner/mock_color_chooser.cc
+++ b/components/test_runner/mock_color_chooser.cc
@@ -17,23 +17,29 @@ MockColorChooser::MockColorChooser(blink::WebColorChooserClient* client,
       delegate_(delegate),
       test_runner_(test_runner),
       weak_factory_(this) {
-  test_runner_->DidOpenChooser();
+  test_runner_->DidOpenChooser(this);
 }
 
 MockColorChooser::~MockColorChooser() {
-  test_runner_->DidCloseChooser();
+  test_runner_->DidCloseChooser(this);
 }
 
 void MockColorChooser::setSelectedColor(const blink::WebColor color) {}
 
 void MockColorChooser::endChooser() {
-  delegate_->PostDelayedTask(base::Bind(&MockColorChooser::InvokeDidEndChooser,
-                                        weak_factory_.GetWeakPtr()),
-                             0);
+  if (delegate_)
+    delegate_->PostDelayedTask(
+        base::Bind(&MockColorChooser::InvokeDidEndChooser,
+                   weak_factory_.GetWeakPtr()),
+        0);
 }
 
 void MockColorChooser::InvokeDidEndChooser() {
   client_->didEndChooser();
 }
 
+void MockColorChooser::SetDelegate(WebTestDelegate* delegate) {
+  delegate_ = delegate;
+}
+
 }  // namespace test_runner