Extract out a common mock task into a module.

The mock tasks are used in multiple unitest modules. It is not extracted out
into a module to enable reuse.

BUG=None
TEST=unit testing for the pipeline stage and pipeline workers.

Change-Id: Iee436266f7ea26c2877377f52208cc5cd31e7815
Reviewed-on: https://gerrit-int.chromium.org/41001
Reviewed-by: Simon Que <sque@google.com>
Reviewed-by: Luis Lozano <llozano@chromium.org>
Commit-Queue: Yuheng Long <yuhenglong@google.com>
Tested-by: Yuheng Long <yuhenglong@google.com>

4 files changed, 80 insertions, 68 deletions

diff --git a/bestflags/mock_task.py b/bestflags/mock_task.py
new file mode 100644
index 00000000..2582f111
--- /dev/null
+++ b/bestflags/mock_task.py
@@ -0,0 +1,64 @@
+"""This module defines the common mock task used by varies unit tests.
+
+Part of the Chrome build flags optimization.
+"""
+
+__author__ = 'yuhenglong@google.com (Yuheng Long)'
+
+# Pick an integer at random.
+POISONPILL = 975
+
+
+class MockTask(object):
+  """This class emulates an actual task.
+
+  It does not do the actual work, but simply returns the result as given when
+  this task is constructed.
+  """
+
+  def __init__(self, stage, identifier, cost=0):
+    """Set up the results for this task.
+
+    Args:
+      stage: the stage of this test is in.
+      identifier: the identifier of this task.
+      cost: the mock cost of this task.
+
+      The _pre_cost field stored the cost. Once this task is performed, i.e., by
+      calling the work method, the _cost field will have this cost. The stage
+      field verifies that the module being tested and the unitest are in the
+      same stage. If the unitest does not care about cost of this task, the cost
+      parameter should be leaved blank.
+    """
+
+    self._identifier = identifier
+    self._pre_cost = cost
+    self._stage = stage
+
+  def __eq__(self, other):
+    if isinstance(other, MockTask):
+      return (self._identifier == other.GetIdentifier(self._stage) and
+              self._cost == other.GetResult(self._stage))
+    return False
+
+  def GetIdentifier(self, stage):
+    assert stage == self._stage
+    return self._identifier
+
+  def SetResult(self, stage, cost):
+    assert stage == self._stage
+    self._cost = cost
+
+  def Work(self, stage):
+    assert stage == self._stage
+    self._cost = self._pre_cost
+
+  def GetResult(self, stage):
+    assert stage == self._stage
+    return self._cost
+
+  def Done(self, stage):
+    """Indicates whether the task has been performed."""
+
+    assert stage == self._stage
+    return '_cost' in self.__dict__
diff --git a/bestflags/pipeline_process.py b/bestflags/pipeline_process.py
index 0999c28b..20f77ab5 100644
--- a/bestflags/pipeline_process.py
+++ b/bestflags/pipeline_process.py
@@ -101,7 +101,7 @@ class PipelineProcess(multiprocessing.Process):
         self._result_queue.put(POISONPILL)
         break
 
-      task_key = task.GetKey(self._stage)
+      task_key = task.GetIdentifier(self._stage)
       if task_key in mycache:
         # The task has been encountered before. It will be sent to the helper
         # module for further processing.
diff --git a/bestflags/pipeline_process_test.py b/bestflags/pipeline_process_test.py
index 086c231b..b5956857 100644
--- a/bestflags/pipeline_process_test.py
+++ b/bestflags/pipeline_process_test.py
@@ -8,10 +8,13 @@ __author__ = 'yuhenglong@google.com (Yuheng Long)'
 import multiprocessing
 import unittest
 
+from mock_task import MockTask
 import pipeline_process
 
 # Pick an integer at random.
 ERROR = -334
+# Pick an integer at random.
+TESTSTAGE = -8
 
 
 def MockHelper(done_dict, helper_queue, _, result_queue):
@@ -27,19 +30,11 @@ def MockHelper(done_dict, helper_queue, _, result_queue):
         # verify that it does not get duplicate "1"s in the test.
         result_queue.put(ERROR)
       else:
-        result_queue.put(('helper', task.GetKey(0)))
+        result_queue.put(('helper', task.GetIdentifier(TESTSTAGE)))
 
 
 def MockWorker(task, _, result_queue):
-  result_queue.put(('worker', task.GetKey(0)))
-
-
-class MockTask(object):
-  def __init__(self, key):
-    self._key = key
-
-  def GetKey(self, _):
-    return self._key
+  result_queue.put(('worker', task.GetIdentifier(TESTSTAGE)))
 
 
 class PipelineProcessTest(unittest.TestCase):
@@ -63,13 +58,13 @@ class PipelineProcessTest(unittest.TestCase):
     inp = manager.Queue()
     output = manager.Queue()
 
-    process = pipeline_process.PipelineProcess(2, 'testing', {}, 'test', inp,
+    process = pipeline_process.PipelineProcess(2, 'testing', {}, TESTSTAGE, inp,
                                                MockHelper, MockWorker, output)
 
     process.start()
-    inp.put(MockTask(1))
-    inp.put(MockTask(1))
-    inp.put(MockTask(2))
+    inp.put(MockTask(TESTSTAGE, 1))
+    inp.put(MockTask(TESTSTAGE, 1))
+    inp.put(MockTask(TESTSTAGE, 2))
     inp.put(pipeline_process.POISONPILL)
     process.join()
 
diff --git a/bestflags/pipeline_worker_test.py b/bestflags/pipeline_worker_test.py
index 0563852e..98d141dd 100644
--- a/bestflags/pipeline_worker_test.py
+++ b/bestflags/pipeline_worker_test.py
@@ -12,11 +12,13 @@ import random
 import sys
 import unittest
 
+from mock_task import MockTask
 import pipeline_process
 import pipeline_worker
 
 
-TESTSTAGE = 0
+# Pick an integer at random.
+TESTSTAGE = -3
 
 
 def MockTaskCostGenerator():
@@ -24,56 +26,7 @@ def MockTaskCostGenerator():
   return random.randint(-sys.maxint - 1, -1)
 
 
-class MockTask(object):
-  """This class emulates an actual task.
-
-  It does not do the actual work, but simply returns the result as given when
-  this task is constructed.
-  """
-
-  def __init__(self, identifier, cost):
-    """Set up the results for this task.
-
-    Args:
-      identifier: the identifier of this task.
-      cost: the mock cost of this task.
-
-      The _pre_cost field stores the cost. Once this task is performed, i.e., by
-      calling the work method , the _cost field will have this cost.
-    """
-
-    self._identifier = identifier
-    self._pre_cost = cost
-
-  def __eq__(self, other):
-    if isinstance(other, MockTask):
-      return self._identifier == other._identifier and self._cost == other._cost
-    return False
-
-  def GetIdentifier(self, stage):
-    assert stage == TESTSTAGE
-    return self._identifier
-
-  def SetResult(self, stage, cost):
-    assert stage == TESTSTAGE
-    self._cost = cost
-
-  def Work(self, stage):
-    assert stage == TESTSTAGE
-    self._cost = self._pre_cost
-
-  def GetResult(self, stage):
-    assert stage == TESTSTAGE
-    return self._cost
-
-  def Done(self, stage):
-    """Indicates whether the task has been performed."""
-
-    assert stage == TESTSTAGE
-    return '_cost' in self.__dict__
-
-
-class AuxiliaryTest(unittest.TestCase):
+class PipelineWorkerTest(unittest.TestCase):
   """This class tests the pipeline_worker functions.
 
   Given the same identifier, the cost should result the same from the
@@ -115,7 +68,7 @@ class AuxiliaryTest(unittest.TestCase):
     # Testing the correctness of having tasks having the same identifier, here
     # 1.
     for result in results:
-      helper_queue.put(MockTask(result, MockTaskCostGenerator()))
+      helper_queue.put(MockTask(TESTSTAGE, result, MockTaskCostGenerator()))
 
     completed_queue.put((2, mock_result[2]))
     completed_queue.put((1, mock_result[1]))
@@ -150,7 +103,7 @@ class AuxiliaryTest(unittest.TestCase):
     mock_tasks = []
 
     for flag, cost in mock_work_tasks.iteritems():
-      mock_tasks.append(MockTask(flag, cost))
+      mock_tasks.append(MockTask(TESTSTAGE, flag, cost))
 
     # Submit the mock tasks to the worker.
     for mock_task in mock_tasks: