1 files changed, 67 insertions, 11 deletions

diff --git a/cc/test/task_graph_runner_test_template.h b/cc/test/task_graph_runner_test_template.h
index d3989b9..d417543 100644
--- a/cc/test/task_graph_runner_test_template.h
+++ b/cc/test/task_graph_runner_test_template.h
@@ -24,17 +24,20 @@ class TaskGraphRunnerTestBase {
              unsigned id,
              unsigned dependent_id,
              unsigned dependent_count,
+             unsigned category,
              unsigned priority)
         : namespace_index(namespace_index),
           id(id),
           dependent_id(dependent_id),
           dependent_count(dependent_count),
+          category(category),
           priority(priority) {}
 
     int namespace_index;
     unsigned id;
     unsigned dependent_id;
     unsigned dependent_count;
+    unsigned category;
     unsigned priority;
   };
 
@@ -126,8 +129,8 @@ TYPED_TEST_P(TaskGraphRunnerTest, Basic) {
     EXPECT_EQ(0u, this->run_task_ids(i).size());
     EXPECT_EQ(0u, this->on_task_completed_ids(i).size());
 
-    this->ScheduleTasks(i,
-                        std::vector<TaskInfo>(1, TaskInfo(i, 0u, 0u, 0u, 0u)));
+    this->ScheduleTasks(
+        i, std::vector<TaskInfo>(1, TaskInfo(i, 0u, 0u, 0u, 0u, 0u)));
   }
 
   for (int i = 0; i < kNamespaceCount; ++i) {
@@ -138,8 +141,8 @@ TYPED_TEST_P(TaskGraphRunnerTest, Basic) {
   }
 
   for (int i = 0; i < kNamespaceCount; ++i)
-    this->ScheduleTasks(i,
-                        std::vector<TaskInfo>(1, TaskInfo(i, 0u, 0u, 1u, 0u)));
+    this->ScheduleTasks(
+        i, std::vector<TaskInfo>(1, TaskInfo(i, 0u, 0u, 1u, 0u, 0u)));
 
   for (int i = 0; i < kNamespaceCount; ++i) {
     this->RunAllTasks(i);
@@ -149,8 +152,8 @@ TYPED_TEST_P(TaskGraphRunnerTest, Basic) {
   }
 
   for (int i = 0; i < kNamespaceCount; ++i)
-    this->ScheduleTasks(i,
-                        std::vector<TaskInfo>(1, TaskInfo(i, 0u, 0u, 2u, 0u)));
+    this->ScheduleTasks(
+        i, std::vector<TaskInfo>(1, TaskInfo(i, 0u, 0u, 2u, 0u, 0u)));
 
   for (int i = 0; i < kNamespaceCount; ++i) {
     this->RunAllTasks(i);
@@ -167,7 +170,7 @@ TYPED_TEST_P(TaskGraphRunnerTest, Dependencies) {
   for (int i = 0; i < kNamespaceCount; ++i) {
     this->ScheduleTasks(i, std::vector<TaskInfo>(1, TaskInfo(i, 0u, 1u,
                                                              1u,  // 1 dependent
-                                                             0u)));
+                                                             0u, 0u)));
   }
 
   for (int i = 0; i < kNamespaceCount; ++i) {
@@ -185,7 +188,7 @@ TYPED_TEST_P(TaskGraphRunnerTest, Dependencies) {
     this->ScheduleTasks(i,
                         std::vector<TaskInfo>(1, TaskInfo(i, 2u, 3u,
                                                           2u,  // 2 dependents
-                                                          0u)));
+                                                          0u, 0u)));
   }
 
   for (int i = 0; i < kNamespaceCount; ++i) {
@@ -201,7 +204,60 @@ TYPED_TEST_P(TaskGraphRunnerTest, Dependencies) {
   }
 }
 
-REGISTER_TYPED_TEST_CASE_P(TaskGraphRunnerTest, Basic, Dependencies);
+TYPED_TEST_P(TaskGraphRunnerTest, Categorys) {
+  const int kNamespaceCount = TaskGraphRunnerTestBase::kNamespaceCount;
+  const unsigned kCategoryCount = 3;
+  using TaskInfo = TaskGraphRunnerTestBase::TaskInfo;
+
+  for (int i = 0; i < kNamespaceCount; ++i) {
+    EXPECT_EQ(0u, this->run_task_ids(i).size());
+    EXPECT_EQ(0u, this->on_task_completed_ids(i).size());
+    std::vector<TaskInfo> tasks;
+    for (unsigned j = 0; j < kCategoryCount; ++j) {
+      tasks.emplace_back(i, 0u, 0u, 0u, j, 0u);
+    }
+    this->ScheduleTasks(i, tasks);
+  }
+
+  for (int i = 0; i < kNamespaceCount; ++i) {
+    this->RunAllTasks(i);
+
+    EXPECT_EQ(kCategoryCount, this->run_task_ids(i).size());
+    EXPECT_EQ(kCategoryCount, this->on_task_completed_ids(i).size());
+  }
+
+  for (int i = 0; i < kNamespaceCount; ++i) {
+    std::vector<TaskInfo> tasks;
+    for (unsigned j = 0; j < kCategoryCount; ++j) {
+      tasks.emplace_back(i, 0u, 0u, 1u, j, 0u);
+    }
+    this->ScheduleTasks(i, tasks);
+  }
+
+  for (int i = 0; i < kNamespaceCount; ++i) {
+    this->RunAllTasks(i);
+
+    EXPECT_EQ(kCategoryCount * 3u, this->run_task_ids(i).size());
+    EXPECT_EQ(kCategoryCount * 2u, this->on_task_completed_ids(i).size());
+  }
+
+  for (int i = 0; i < kNamespaceCount; ++i) {
+    std::vector<TaskInfo> tasks;
+    for (unsigned j = 0; j < kCategoryCount; ++j) {
+      tasks.emplace_back(i, 0u, 0u, 2u, j, 0u);
+    }
+    this->ScheduleTasks(i, tasks);
+  }
+
+  for (int i = 0; i < kNamespaceCount; ++i) {
+    this->RunAllTasks(i);
+
+    EXPECT_EQ(kCategoryCount * 6u, this->run_task_ids(i).size());
+    EXPECT_EQ(kCategoryCount * 3u, this->on_task_completed_ids(i).size());
+  }
+}
+
+REGISTER_TYPED_TEST_CASE_P(TaskGraphRunnerTest, Basic, Dependencies, Categorys);
 
 template <typename TaskRunnerTestDelegate>
 using SingleThreadTaskGraphRunnerTest =
@@ -215,8 +271,8 @@ TYPED_TEST_P(SingleThreadTaskGraphRunnerTest, Priority) {
 
   for (int i = 0; i < kNamespaceCount; ++i) {
     TaskInfo tasks[] = {
-        TaskInfo(i, 0u, 2u, 1u, 1u),  // Priority 1
-        TaskInfo(i, 1u, 3u, 1u, 0u)   // Priority 0
+        TaskInfo(i, 0u, 2u, 1u, 0u, 1u),  // Priority 1
+        TaskInfo(i, 1u, 3u, 1u, 0u, 0u)   // Priority 0
     };
     this->ScheduleTasks(i,
                         std::vector<TaskInfo>(tasks, tasks + arraysize(tasks)));