diff --git a/bestflags/generation_test.py b/bestflags/generation_test.py
index 748a988..1249f7d 100644
--- a/bestflags/generation_test.py
+++ b/bestflags/generation_test.py
@@ -13,7 +13,7 @@
 import unittest
 
 from generation import Generation
-from mock_task import MockTask
+from mock_task import IdentifierMockTask
 
 
 # Pick an integer at random.
@@ -27,24 +27,6 @@
 STRIDE = 7
 
 
-class GenerationMockTask(MockTask):
-  """This class defines the mock task to test the Generation class.
-
-  The task instances will be inserted into a set. Therefore the hash and the
-  equal methods are overridden. The generation class considers the identifier to
-  set the cost of the task in a set, thus the identifier is used in the
-  overriding methods.
-  """
-
-  def __hash__(self):
-    return self._identifier
-
-  def __eq__(self, other):
-    if isinstance(other, MockTask):
-      return self._identifier == other.GetIdentifier(self._stage)
-    return False
-
-
 class GenerationTest(unittest.TestCase):
   """This class test the Generation class.
 
@@ -66,7 +48,7 @@
     testing_tasks = range(NUMTASKS)
 
     # The tasks for the generation to be tested.
-    generation_tasks = [GenerationMockTask(TESTSTAGE, t) for t in testing_tasks]
+    generation_tasks = [IdentifierMockTask(TESTSTAGE, t) for t in testing_tasks]
 
     gen = Generation(set(generation_tasks), None)
 
@@ -81,7 +63,7 @@
 
       # Mark a task as done by calling the UpdateTask method of the generation.
       # Send the generation the task as well as its results.
-      gen.UpdateTask(GenerationMockTask(TESTSTAGE, testing_task))
+      gen.UpdateTask(IdentifierMockTask(TESTSTAGE, testing_task))
 
     # The Done method should return true after all the tasks in the permuted
     # list is set.