Notify DataUseTabModel of navigations and tab closures

chrome/browser/android/data_usage/data_use_tab_model_unittest.cc

Index: chrome/browser/android/data_usage/data_use_tab_model_unittest.cc
diff --git a/chrome/browser/android/data_usage/data_use_tab_model_unittest.cc b/chrome/browser/android/data_usage/data_use_tab_model_unittest.cc
index 70777f6451d0c0346df7c7921d035ff7f3e17786..f7ebe7c477ab447ed39afbf4577e9314bc73b5ba 100644
--- a/chrome/browser/android/data_usage/data_use_tab_model_unittest.cc
+++ b/chrome/browser/android/data_usage/data_use_tab_model_unittest.cc
@@ -8,17 +8,21 @@
 
 #include <string>
 
+#include "base/memory/weak_ptr.h"
 #include "base/message_loop/message_loop.h"
+#include "base/run_loop.h"
 #include "base/single_thread_task_runner.h"
 #include "base/strings/stringprintf.h"
 #include "base/test/histogram_tester.h"
 #include "base/time/time.h"
 #include "chrome/browser/android/data_usage/external_data_use_observer.h"
 #include "chrome/browser/android/data_usage/tab_data_use_entry.h"
+#include "chrome/test/base/chrome_render_view_host_test_harness.h"
 #include "components/data_usage/core/data_use.h"
 #include "components/data_usage/core/data_use_aggregator.h"
 #include "components/data_usage/core/data_use_amortizer.h"
 #include "components/data_usage/core/data_use_annotator.h"
+#include "content/public/browser/browser_thread.h"
 #include "net/base/network_change_notifier.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -71,20 +75,27 @@ class DataUseTabModelNowTest : public DataUseTabModel {
   base::TimeDelta now_offset_;
 };
 
-class DataUseTabModelTest : public testing::Test {
+class DataUseTabModelTest : public ChromeRenderViewHostTestHarness {
  public:
   DataUseTabModelTest() {}
 
  protected:
   void SetUp() override {
+    ChromeRenderViewHostTestHarness::SetUp();
+
     data_use_aggregator_.reset(new data_usage::DataUseAggregator(
         scoped_ptr<data_usage::DataUseAnnotator>(),
         scoped_ptr<data_usage::DataUseAmortizer>()));
     data_use_observer_.reset(new ExternalDataUseObserver(
-        data_use_aggregator_.get(), message_loop_.task_runner().get(),
-        message_loop_.task_runner().get()));
+        data_use_aggregator_.get(),
+        content::BrowserThread::GetMessageLoopProxyForThread(
+            content::BrowserThread::IO),
+        content::BrowserThread::GetMessageLoopProxyForThread(
+            content::BrowserThread::UI)));
     data_use_tab_model_ = new DataUseTabModelNowTest(
-        data_use_observer_.get(), message_loop_.task_runner().get());
+        data_use_observer_.get(),
+        content::BrowserThread::GetMessageLoopProxyForThread(
+            content::BrowserThread::UI));
 
     // |data_use_tab_model_| will be owned by |data_use_observer_|.
     data_use_observer_->data_use_tab_model_.reset(data_use_tab_model_);
@@ -174,14 +185,25 @@ class DataUseTabModelTest : public testing::Test {
   // Pointer to the tab model within and owned by ExternalDataUseObserver.
   DataUseTabModelNowTest* data_use_tab_model_;
 
-  base::MessageLoop message_loop_;
+ private:
+  DISALLOW_COPY_AND_ASSIGN(DataUseTabModelTest);
 };
 
 // Mock observer to track the calls to start and end tracking events.
 class MockTabDataUseObserver : public DataUseTabModel::TabDataUseObserver {
  public:
+  MockTabDataUseObserver() : weak_ptr_factory_(this) {}
   MOCK_METHOD1(NotifyTrackingStarting, void(int32_t tab_id));
   MOCK_METHOD1(NotifyTrackingEnding, void(int32_t tab_id));
+
+  base::WeakPtr<MockTabDataUseObserver> GetWeakPtr() {
+    return weak_ptr_factory_.GetWeakPtr();
+  }
+
+ private:
+  base::WeakPtrFactory<MockTabDataUseObserver> weak_ptr_factory_;
+
+  DISALLOW_COPY_AND_ASSIGN(MockTabDataUseObserver);
 };
 
 // Starts and ends tracking a single tab and checks if its label is returned
@@ -249,21 +271,27 @@ TEST_F(DataUseTabModelTest, ObserverStartEndEvents) {
   EXPECT_CALL(mock_observer, NotifyTrackingStarting(kTabID1)).Times(1);
   EXPECT_CALL(mock_observer, NotifyTrackingEnding(kTabID1)).Times(1);
 
-  data_use_tab_model_->AddObserver(&mock_observer);
+  data_use_tab_model_->AddObserver(&mock_observer, mock_observer.GetWeakPtr());
+  EXPECT_EQ(1U, data_use_tab_model_->observers_.size());
   StartTrackingDataUse(kTabID1, kTestLabel1);
   EndTrackingDataUse(kTabID1);
 
-  message_loop_.RunUntilIdle();
+  base::RunLoop().RunUntilIdle();
 }
 
 // Checks that multiple mock observers receive start and end tracking events for
 // multiple tabs.
 TEST_F(DataUseTabModelTest, MultipleObserverMultipleStartEndEvents) {
   MockTabDataUseObserver mock_observers[kMaxMockObservers];
+  size_t expected_observer_count = 0;
+  EXPECT_EQ(expected_observer_count, data_use_tab_model_->observers_.size());
 
   for (auto& mock_observer : mock_observers) {
     // Add the observer.
-    data_use_tab_model_->AddObserver(&mock_observer);
+    data_use_tab_model_->AddObserver(&mock_observer,
+                                     mock_observer.GetWeakPtr());
+    ++expected_observer_count;
+    EXPECT_EQ(expected_observer_count, data_use_tab_model_->observers_.size());
 
     // Expect start and end events for tab ids 1-3.
     EXPECT_CALL(mock_observer, NotifyTrackingStarting(kTabID1)).Times(1);
@@ -282,7 +310,15 @@ TEST_F(DataUseTabModelTest, MultipleObserverMultipleStartEndEvents) {
   EndTrackingDataUse(kTabID2);
   EndTrackingDataUse(kTabID3);
 
-  message_loop_.RunUntilIdle();
+  base::RunLoop().RunUntilIdle();
+
+  for (auto& mock_observer : mock_observers) {
+    // Remove the observer.
+    data_use_tab_model_->RemoveObserver(&mock_observer);
+    --expected_observer_count;
+    EXPECT_EQ(expected_observer_count, data_use_tab_model_->observers_.size());
+  }
+  EXPECT_EQ(0U, data_use_tab_model_->observers_.size());
 }
 
 // Checks that the observer is not notified of start and end events after
@@ -294,11 +330,13 @@ TEST_F(DataUseTabModelTest, ObserverNotNotifiedAfterRemove) {
   EXPECT_CALL(mock_observer, NotifyTrackingStarting(kTabID1)).Times(1);
   EXPECT_CALL(mock_observer, NotifyTrackingEnding(kTabID1)).Times(1);
 
-  data_use_tab_model_->AddObserver(&mock_observer);
+  EXPECT_EQ(0U, data_use_tab_model_->observers_.size());
+  data_use_tab_model_->AddObserver(&mock_observer, mock_observer.GetWeakPtr());
+  EXPECT_EQ(1U, data_use_tab_model_->observers_.size());
   StartTrackingDataUse(kTabID1, kTestLabel1);
   EndTrackingDataUse(kTabID1);
 
-  message_loop_.RunUntilIdle();
+  base::RunLoop().RunUntilIdle();
   testing::Mock::VerifyAndClear(&mock_observer);
 
   // Observer should not be notified after RemoveObserver.
@@ -306,10 +344,11 @@ TEST_F(DataUseTabModelTest, ObserverNotNotifiedAfterRemove) {
   EXPECT_CALL(mock_observer, NotifyTrackingEnding(kTabID1)).Times(0);
 
   data_use_tab_model_->RemoveObserver(&mock_observer);
+  EXPECT_EQ(0U, data_use_tab_model_->observers_.size());
   StartTrackingDataUse(kTabID1, kTestLabel1);
   EndTrackingDataUse(kTabID1);
 
-  message_loop_.RunUntilIdle();
+  base::RunLoop().RunUntilIdle();
 }
 
 // Checks that tab close event updates the close time of the tab entry.
@@ -348,7 +387,7 @@ TEST_F(DataUseTabModelTest, OnTrackingLabelRemoved) {
   StartTrackingDataUse(kTabID1, kTestLabel1);
   StartTrackingDataUse(kTabID2, kTestLabel2);
   StartTrackingDataUse(kTabID3, kTestLabel3);
-  data_use_tab_model_->AddObserver(&mock_observer);
+  data_use_tab_model_->AddObserver(&mock_observer, mock_observer.GetWeakPtr());
   ExpectTabEntrySize(TabEntrySize::THREE);
 
   EXPECT_TRUE(IsTrackingDataUse(kTabID1));
@@ -359,7 +398,7 @@ TEST_F(DataUseTabModelTest, OnTrackingLabelRemoved) {
   EXPECT_CALL(mock_observer, NotifyTrackingEnding(kTabID2)).Times(1);
 
   data_use_tab_model_->OnTrackingLabelRemoved(kTestLabel2);
-  message_loop_.RunUntilIdle();
+  base::RunLoop().RunUntilIdle();
 
   EXPECT_TRUE(IsTrackingDataUse(kTabID1));
   EXPECT_FALSE(IsTrackingDataUse(kTabID2));
@@ -368,7 +407,7 @@ TEST_F(DataUseTabModelTest, OnTrackingLabelRemoved) {
   EXPECT_CALL(mock_observer, NotifyTrackingEnding(kTabID3)).Times(1);
 
   data_use_tab_model_->OnTrackingLabelRemoved(kTestLabel3);
-  message_loop_.RunUntilIdle();
+  base::RunLoop().RunUntilIdle();
 
   EXPECT_TRUE(IsTrackingDataUse(kTabID1));
   EXPECT_FALSE(IsTrackingDataUse(kTabID2));
@@ -558,11 +597,6 @@ TEST_F(DataUseTabModelTest, NavigationEnterAndExitEvent) {
       GURL("http://foo.com/"), std::string());
   ExpectTabEntrySize(TabEntrySize::ONE);
   EXPECT_TRUE(IsTrackingDataUse(kTabID1));
-
-  data_use_tab_model_->OnNavigationEvent(
-      kTabID1, DataUseTabModel::TRANSITION_FROM_NAVSUGGEST, GURL(),
-      std::string());
-  EXPECT_FALSE(IsTrackingDataUse(kTabID1));
 }
 
 // Tests that any of the Enter transition events start the tracking.
@@ -611,9 +645,6 @@ TEST_F(DataUseTabModelTest, AllNavigationEnterEvents) {
 // Tests that any of the Exit transition events end the tracking.
 TEST_F(DataUseTabModelTest, AllNavigationExitEvents) {
   DataUseTabModel::TransitionType all_exit_transitions[] = {
-      DataUseTabModel::TRANSITION_TO_EXTERNAL_APP,
-      DataUseTabModel::TRANSITION_FROM_NAVSUGGEST,
-      DataUseTabModel::TRANSITION_OMNIBOX_NAVIGATION,
       DataUseTabModel::TRANSITION_BOOKMARK,
       DataUseTabModel::TRANSITION_HISTORY_ITEM};
   std::vector<std::string> app_package_names, domain_regexes, labels;