[Paint Preview] Correct error handling

Fix an issue where TabService attempted to proceed to write the proto
even if there was an error.

Change-Id: I5734921cd77baa026cc47e47d2656b87f672aa94
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2095012Reviewed-by: Mehran Mahmoudi <mahmoudi@chromium.org>
Commit-Queue: Calder Kitagawa <ckitagawa@chromium.org>
Cr-Commit-Position: refs/heads/master@{#748685}

chrome/browser/paint_preview/services/paint_preview_tab_service.cc

@@ -77,7 +77,10 @@ void PaintPreviewTabService::OnCaptured(
     FinishedCallback callback,
     PaintPreviewBaseService::CaptureStatus status,
     std::unique_ptr<PaintPreviewProto> proto) {
-  auto result = base::MakeRefCounted<base::RefCountedData<bool>>(false);
+  if (status != PaintPreviewBaseService::CaptureStatus::kOk || !proto) {
+    std::move(callback).Run(Status::kCaptureFailed);
+    return;
+  }
   auto file_manager = GetFileManager();
   GetTaskRunner()->PostTaskAndReplyWithResult(
       FROM_HERE,

chrome/browser/paint_preview/services/paint_preview_tab_service_unittest.cc

@@ -138,6 +138,44 @@ TEST_F(PaintPreviewTabServiceTest, CaptureTab) {
   content::RunAllTasksUntilIdle();
 }
 
+TEST_F(PaintPreviewTabServiceTest, CaptureTabFailed) {
+  content::NavigationSimulator::NavigateAndCommitFromBrowser(
+      web_contents(), GURL("http://www.example.com"));
+  const int kTabId = 1U;
+
+  MockPaintPreviewRecorder recorder;
+  recorder.SetResponse(mojom::PaintPreviewStatus::kFailed);
+  OverrideInterface(&recorder);
+
+  auto* service = GetService();
+  base::RunLoop loop;
+  service->CaptureTab(
+      kTabId, web_contents(),
+      base::BindOnce(
+          [](base::OnceClosure quit, PaintPreviewTabService::Status status) {
+            EXPECT_EQ(status, PaintPreviewTabService::Status::kCaptureFailed);
+            std::move(quit).Run();
+          },
+          loop.QuitClosure()));
+  loop.Run();
+
+  auto file_manager = service->GetFileManager();
+  auto key = file_manager->CreateKey(kTabId);
+  service->GetTaskRunner()->PostTaskAndReplyWithResult(
+      FROM_HERE,
+      base::BindOnce(&FileManager::DirectoryExists, file_manager, key),
+      base::BindOnce([](bool exists) { EXPECT_TRUE(exists); }));
+  content::RunAllTasksUntilIdle();
+
+  service->TabClosed(kTabId);
+  content::RunAllTasksUntilIdle();
+  service->GetTaskRunner()->PostTaskAndReplyWithResult(
+      FROM_HERE,
+      base::BindOnce(&FileManager::DirectoryExists, file_manager, key),
+      base::BindOnce([](bool exists) { EXPECT_FALSE(exists); }));
+  content::RunAllTasksUntilIdle();
+}
+
 TEST_F(PaintPreviewTabServiceTest, CaptureTabTwice) {
   content::NavigationSimulator::NavigateAndCommitFromBrowser(
       web_contents(), GURL("http://www.example.com"));

components/paint_preview/browser/paint_preview_base_service.cc

@@ -114,7 +114,7 @@ void PaintPreviewBaseService::OnCaptured(
     base::UnguessableToken guid,
     mojom::PaintPreviewStatus status,
     std::unique_ptr<PaintPreviewProto> proto) {
-  if (status != mojom::PaintPreviewStatus::kOk) {
+  if (status != mojom::PaintPreviewStatus::kOk || !proto) {
     DVLOG(1) << "ERROR: Paint Preview failed to capture for document "
              << guid.ToString() << " with error " << status;
     std::move(callback).Run(kCaptureFailed, nullptr);