class SessionContainer(object): """ Session container object. """ def __init__(self): self._items = dict() self._queue = PriorityQueue() def add(self, session): """ Add session to the container. @param session: Session object """ self._items[session.session_id] = session if session.expiry is not None: self._queue.push(session) def get(self, session_id): """ Return session object or None if it is not available @param session_id: Session identifier """ return self._items.get(session_id, None) def remove(self, session_id): """ Remove session object from the container @param session_id: Session identifier """ session = self._items.get(session_id, None) if session is not None: session.promoted = -1 session.on_delete(True) del self._items[session_id] return True return False def expire(self, current_time=None): """ Expire any old entries @param current_time: Optional time to be used to clean up queue (can be used in unit tests) """ if self._queue.is_empty(): return if current_time is None: current_time = time.time() while not self._queue.is_empty(): # Get top most item top = self._queue.peek() # Early exit if item was not promoted and its expiration time # is greater than now. if top.promoted is None and top.expiry_date > current_time: break # Pop item from the stack top = self._queue.pop() need_reschedule = (top.promoted is not None and top.promoted > current_time) # Give chance to reschedule if not need_reschedule: top.promoted = None top.on_delete(False) need_reschedule = (top.promoted is not None and top.promoted > current_time) # If item is promoted and expiration time somewhere in future # just reschedule it if need_reschedule: top.expiry_date = top.promoted top.promoted = None self._queue.push(top) else: del self._items[top.session_id]
class PriorityQueueTest(unittest.TestCase): def setUp(self): self.q = PriorityQueue() def test_is_empty(self): self.assertTrue(self.q.is_empty()) self.q.push(1) self.assertFalse(self.q.is_empty()) self.q.pop() self.assertTrue(self.q.is_empty()) def test_push(self): self.q.push(1) self.assertEquals(self.q.pop(), 1) def test_pop_from_empty_queue_raises_indexerror(self): self.assertRaises(IndexError, self.q.pop) def test_peek_doesnt_pop(self): self.q.push(1) self.assertEquals(self.q.peek(), 1) self.assertFalse(self.q.is_empty()) def test_insertion_order_is_preserved_for_same_priority_elements(self): class El(object): def __init__(self, id, val): self.id = id self.val = val def __cmp__(self, other): return cmp(self.val, other.val) self.q.push(El('other', 2)) self.q.push(El('first', 1)) self.q.push(El('second', 1)) self.q.push(El('third', 1)) self.assertEquals(self.q.pop().id, 'first') self.assertEquals(self.q.pop().id, 'second') self.assertEquals(self.q.pop().id, 'third') self.assertEquals(self.q.pop().id, 'other') def test_contains(self): self.assertFalse(1 in self.q) self.q.push(1) self.assertTrue(1 in self.q) self.q.pop() self.assertFalse(1 in self.q) def test_len(self): self.assertEquals(len(self.q), 0) self.q.push(1) self.assertEquals(len(self.q), 1) self.q.pop() self.assertEquals(len(self.q), 0)