Fix bug with TaskQueueSelector and blocked queues

components/scheduler/base/task_queue_selector_unittest.cc

 // Copyright 2015 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 "components/scheduler/base/task_queue_selector.h"
 
 #include <stddef.h>
 
 #include "base/bind.h"
 #include "base/macros.h"
@@ skipping 108 matching lines @@
     for (size_t i = 0; i < kTaskQueueCount; i++) {
       EXPECT_EQ(TaskQueue::NORMAL_PRIORITY, task_queues_[i]->GetQueuePriority())
           << i;
       queue_to_index_map_.insert(std::make_pair(task_queues_[i].get(), i));
     }
   }
 
   void TearDown() final {
     for (scoped_refptr<TaskQueueImpl>& task_queue : task_queues_) {
       task_queue->UnregisterTaskQueue();
+      // Note since this test doesn't have a TaskQueueManager we need to
+      // manually remove |task_queue| from the |selector_|.  Normally
+      // UnregisterTaskQueue would do that.
+      selector_.RemoveQueue(task_queue.get());
     }
   }
 
   scoped_refptr<TaskQueueImpl> NewTaskQueueWithBlockReporting() {
     return make_scoped_refptr(new TaskQueueImpl(
         nullptr, virtual_time_domain_.get(),
         TaskQueue::Spec("test queue").SetShouldReportWhenExecutionBlocked(true),
         "test", "test"));
   }
 
@@ skipping 41 matching lines @@
 TEST_F(TaskQueueSelectorTest, TestObserverWithEnabledQueue) {
   DisableQueue(task_queues_[1].get());
   MockObserver mock_observer;
   selector_.SetTaskQueueSelectorObserver(&mock_observer);
   EXPECT_CALL(mock_observer, OnTaskQueueEnabled(_)).Times(1);
   EnableQueue(task_queues_[1].get());
 }
 
 TEST_F(TaskQueueSelectorTest,
        TestObserverWithSetQueuePriorityAndQueueAlreadyEnabled) {
-  selector_.SetQueuePriority(task_queues_[1].get(), TaskQueue::NORMAL_PRIORITY);
+  selector_.SetQueuePriority(task_queues_[1].get(), TaskQueue::HIGH_PRIORITY);
   MockObserver mock_observer;
   selector_.SetTaskQueueSelectorObserver(&mock_observer);
   EXPECT_CALL(mock_observer, OnTaskQueueEnabled(_)).Times(0);
   selector_.SetQueuePriority(task_queues_[1].get(), TaskQueue::NORMAL_PRIORITY);
 }
 
 TEST_F(TaskQueueSelectorTest, TestDisableEnable) {
   MockObserver mock_observer;
   selector_.SetTaskQueueSelectorObserver(&mock_observer);
 
@@ skipping 86 matching lines @@
   EXPECT_GT(counts[0], 0ul);        // Check high doesn't starve normal.
   EXPECT_GT(counts[2], counts[0]);  // Check high gets more chance to run.
   EXPECT_EQ(0ul, counts[1]);        // Check best effort is starved.
 }
 
 TEST_F(TaskQueueSelectorTest, TestBestEffortGetsStarved) {
   size_t queue_order[] = {0, 1};
   PushTasks(queue_order, 2);
   selector_.SetQueuePriority(task_queues_[0].get(),
                              TaskQueue::BEST_EFFORT_PRIORITY);
-  selector_.SetQueuePriority(task_queues_[1].get(), TaskQueue::NORMAL_PRIORITY);
+  EXPECT_EQ(TaskQueue::NORMAL_PRIORITY, task_queues_[1]->GetQueuePriority());
   WorkQueue* chosen_work_queue = nullptr;
   for (int i = 0; i < 100; i++) {
     EXPECT_TRUE(selector_.SelectWorkQueueToService(&chosen_work_queue));
     EXPECT_EQ(task_queues_[1].get(), chosen_work_queue->task_queue());
     // Don't remove task from queue to simulate all queues still being full.
   }
   selector_.SetQueuePriority(task_queues_[1].get(), TaskQueue::HIGH_PRIORITY);
   for (int i = 0; i < 100; i++) {
     EXPECT_TRUE(selector_.SelectWorkQueueToService(&chosen_work_queue));
     EXPECT_EQ(task_queues_[1].get(), chosen_work_queue->task_queue());
@@ skipping 75 matching lines @@
 
   task_queue->immediate_work_queue()->PushAndSetEnqueueOrder(
       TaskQueueImpl::Task(FROM_HERE, test_closure_, base::TimeTicks(), 0, true),
       0);
 
   WorkQueue* chosen_work_queue;
   EXPECT_CALL(mock_observer, OnTriedToSelectBlockedWorkQueue(_)).Times(1);
   EXPECT_FALSE(selector.SelectWorkQueueToService(&chosen_work_queue));
 
   task_queue->UnregisterTaskQueue();
+  selector.RemoveQueue(task_queue.get());
 }
 
 TEST_F(TaskQueueSelectorTest, TestObserverWithTwoBlockedQueues) {
   TaskQueueSelectorForTest selector;
   MockObserver mock_observer;
   selector.SetTaskQueueSelectorObserver(&mock_observer);
 
   scoped_refptr<TaskQueueImpl> task_queue(NewTaskQueueWithBlockReporting());
   scoped_refptr<TaskQueueImpl> task_queue2(NewTaskQueueWithBlockReporting());
   selector.AddQueue(task_queue.get());
@@ skipping 12 matching lines @@
       0);
 
   // Should still only see one call to OnTriedToSelectBlockedWorkQueue.
   WorkQueue* chosen_work_queue;
   EXPECT_CALL(mock_observer, OnTriedToSelectBlockedWorkQueue(_)).Times(1);
   EXPECT_FALSE(selector.SelectWorkQueueToService(&chosen_work_queue));
   testing::Mock::VerifyAndClearExpectations(&mock_observer);
 
   // Removing the second queue and selecting again should result in another
   // notification.
+  task_queue->UnregisterTaskQueue();
   selector.RemoveQueue(task_queue.get());
   EXPECT_CALL(mock_observer, OnTriedToSelectBlockedWorkQueue(_)).Times(1);
   EXPECT_FALSE(selector.SelectWorkQueueToService(&chosen_work_queue));
 
-  task_queue->UnregisterTaskQueue();
   task_queue2->UnregisterTaskQueue();
+  selector.RemoveQueue(task_queue2.get());
 }
 
 struct ChooseOldestWithPriorityTestParam {
   int delayed_task_enqueue_order;
   int immediate_task_enqueue_order;
   int immediate_starvation_count;
   const char* expected_work_queue_name;
   bool expected_did_starve_immediate_queue;
 };
 
@@ skipping 36 matching lines @@
   EXPECT_EQ(chose_delayed_over_immediate,
             GetParam().expected_did_starve_immediate_queue);
 }
 
 INSTANTIATE_TEST_CASE_P(ChooseOldestWithPriorityTest,
                         ChooseOldestWithPriorityTest,
                         testing::ValuesIn(kChooseOldestWithPriorityTestCases));
 
 }  // namespace internal
 }  // namespace scheduler