4 files changed, 36 insertions, 10 deletions
diff --git a/Tools/QueueStatusServer/model/activeworkitems.py b/Tools/QueueStatusServer/model/activeworkitems.py
index ab5d7a6..023258a 100644
--- a/Tools/QueueStatusServer/model/activeworkitems.py
+++ b/Tools/QueueStatusServer/model/activeworkitems.py
@@ -61,6 +61,14 @@ class ActiveWorkItems(db.Model, QueuePropertyMixin):
         nonexpired_pairs = [pair for pair in self._item_time_pairs() if pair[0] != item_id]
         self._set_item_time_pairs(nonexpired_pairs)
 
+    @classmethod
+    def key_for_queue(cls, queue_name):
+        return "active-work-items-%s" % (queue_name)
+
+    @classmethod
+    def lookup_by_queue(cls, queue_name):
+        return cls.get_or_insert(key_name=cls.key_for_queue(queue_name), queue_name=queue_name)
+
     @staticmethod
     def _expire_item(key, item_id):
         active_work_items = db.get(key)
diff --git a/Tools/QueueStatusServer/model/activeworkitems_unitest.py b/Tools/QueueStatusServer/model/activeworkitems_unittest.py
index 6d915a1..b8c620a 100644
--- a/Tools/QueueStatusServer/model/activeworkitems_unitest.py
+++ b/Tools/QueueStatusServer/model/activeworkitems_unittest.py
@@ -28,15 +28,27 @@
 
 import unittest
 from datetime import datetime
+from google.appengine.ext import testbed
 
 from model.activeworkitems import ActiveWorkItems
 
 
 class ActiveWorkItemsTest(unittest.TestCase):
+
+    def setUp(self):
+        self.testbed = testbed.Testbed()
+        self.testbed.activate()
+        self.testbed.init_datastore_v3_stub()
+        self.testbed.init_memcache_stub()
+
+    def tearDown(self):
+        self.testbed.deactivate()
+
     def test_basic(self):
-        items = ActiveWorkItems()
+        items = ActiveWorkItems.lookup_by_queue("test-queue")
         queued_items = [1, 2]
-        time = datetime.now()
+        # db.Model only stores dates to second resolution, so we use an explicit datetime without milliseconds.
+        time = datetime(2011, 4, 18, 18, 50, 44)
         self.assertEqual(items.next_item(queued_items, time), 1)
         self.assertEqual(items.next_item([1], time), None)
         self.assertEqual(items.next_item([], time), None)
@@ -45,8 +57,8 @@ class ActiveWorkItemsTest(unittest.TestCase):
         self.assertEqual(items.time_for_item(2), None)
 
         items.expire_item(1)
+        # expire_item uses a transaction so it doesn't take effect on the current object.
+        self.assertEqual(items.time_for_item(1), time)
+        # If we look up the saved object, we see it's been updated.
+        items = ActiveWorkItems.lookup_by_queue("test-queue")
         self.assertEqual(items.time_for_item(1), None)
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/Tools/QueueStatusServer/model/queues.py b/Tools/QueueStatusServer/model/queues.py
index 1d46f89..bd5e7ed 100644
--- a/Tools/QueueStatusServer/model/queues.py
+++ b/Tools/QueueStatusServer/model/queues.py
@@ -70,13 +70,11 @@ class Queue(object):
         return self._name
 
     def work_items(self):
-        key_name = "work-items-%s" % (self._name)
-        return WorkItems.get_or_insert(key_name=key_name, queue_name=self._name)
+        return WorkItems.lookup_by_queue(self._name)
 
     # FIXME: active_work_items is a bad name for this lock-table.
     def active_work_items(self):
-        key_name = "active-work-items-%s" % (self._name)
-        return ActiveWorkItems.get_or_insert(key_name=key_name, queue_name=self._name)
+        return ActiveWorkItems.lookup_by_queue(self._name)
 
     def _caplitalize_after_dash(self, string):
         return "-".join([word[0].upper() + word[1:] for word in string.split("-")])
diff --git a/Tools/QueueStatusServer/model/workitems.py b/Tools/QueueStatusServer/model/workitems.py
index 772fc39..b74f5f5 100644
--- a/Tools/QueueStatusServer/model/workitems.py
+++ b/Tools/QueueStatusServer/model/workitems.py
@@ -36,6 +36,14 @@ class WorkItems(db.Model, QueuePropertyMixin):
     item_ids = db.ListProperty(int)
     date = db.DateTimeProperty(auto_now_add=True)
 
+    @classmethod
+    def key_for_queue(cls, queue_name):
+        return "work-items-%s" % (queue_name)
+
+    @classmethod
+    def lookup_by_queue(cls, queue_name):
+        return cls.get_or_insert(key_name=cls.key_for_queue(queue_name), queue_name=queue_name)
+
     def display_position_for_attachment(self, attachment_id):
         """Returns a 1-based index corresponding to the position
         of the attachment_id in the queue.  If the attachment is