Fixed multiple issues in cached disk dict

w3af/core/data/db/cached_disk_dict.py

@@ -133,6 +133,11 @@ def _move_key_to_memory_if_needed(self, key, keys_for_memory):
         :return: The name of the key that was moved to memory, or None if
                  all the keys are still on disk.
         """
+        # The key is already in memory, nothing to do here
+        if key in self._in_memory:
+            return
+
+        # The key must not be in memory, nothing to do here
         if key not in keys_for_memory:
             return
 
@@ -145,8 +150,12 @@ def _move_key_to_memory_if_needed(self, key, keys_for_memory):
             return key
 
     def __setitem__(self, key, value):
-        if len(self._in_memory) < self._max_in_memory:
+        if key in self._in_memory:
             self._in_memory[key] = value
+
+        elif len(self._in_memory) < self._max_in_memory:
+            self._in_memory[key] = value
+
         else:
             self._disk_dict[key] = value
 
@@ -191,6 +200,9 @@ def __iter__(self):
         for key in self._disk_dict:
             yield key
 
+    def iteritems(self):
+        for key, value in self._in_memory.iteritems():
+            yield key, value
 
-def sort_by_value(a, b):
-    return cmp(b[1], a[1])
+        for key, value in self._disk_dict.iteritems():
+            yield key, value

w3af/core/data/db/tests/test_cached_disk_dict.py

@@ -32,42 +32,50 @@ def tearDown(self):
         self.cdd.cleanup()
 
     def test_simple_all_in_memory(self):
-        self.cdd[1] = None
+        self.cdd[1] = 6789
         self.cdd[2] = None
         self.cdd[3] = None
 
-        self.assertEqual(len(self.cdd._in_memory), 3)
         self.assertEqual(len(self.cdd._disk_dict), 0)
+        self.assertEqual(len(self.cdd._in_memory), 3)
+        self.assertEqual(self.cdd._access_count, {1: 1, 2: 1, 3: 1})
+        self.assertEqual(self.cdd._in_memory, {1: 6789, 2: None, 3: None})
+
+        self.assertEqual(self.cdd[1], 6789)
 
     def test_one_in_disk(self):
-        self.cdd[1] = None
+        self.cdd[1] = 6789
         self.cdd[2] = None
         self.cdd[3] = None
 
-        self.cdd[1] = None
+        self.cdd[1] = 6789
         self.cdd[2] = None
         self.cdd[3] = None
 
-        self.cdd[4] = None
+        self.cdd[4] = 9876
 
+        self.assertEqual(self.cdd._access_count, {1: 2, 2: 2, 3: 2, 4: 1})
         self.assertEqual(self.cdd._in_memory.keys(), [1, 2, 3])
         self.assertEqual(self.cdd._disk_dict.keys(), [4])
+        self.assertEqual(self.cdd[1], 6789)
+        self.assertEqual(self.cdd[4], 9876)
 
     def test_one_in_disk_then_moves_to_memory(self):
-        self.cdd[1] = None
+        self.cdd[1] = 6789
         self.cdd[2] = None
         self.cdd[3] = None
 
-        self.cdd[1] = None
+        self.cdd[1] = 6789
         self.cdd[2] = None
         self.cdd[3] = None
 
-        self.cdd[4] = None
+        self.cdd[4] = 9876
 
         self.assertEqual(self.cdd._in_memory.keys(), [1, 2, 3])
         self.assertEqual(self.cdd._disk_dict.keys(), [4])
+        self.assertEqual(self.cdd[1], 6789)
         self.assertEqual(self.cdd._access_count,
-                         {1: 2,
+                         {1: 3,
                           2: 2,
                           3: 2,
                           4: 1})
@@ -77,24 +85,25 @@ def test_one_in_disk_then_moves_to_memory(self):
         self.cdd[4]
 
         self.assertEqual(self.cdd._access_count,
-                         {1: 2,
+                         {1: 3,
                           2: 2,
                           3: 2,
                           4: 4})
 
         self.assertEqual(self.cdd._in_memory.keys(), [1, 2, 4])
         self.assertEqual(self.cdd._disk_dict.keys(), [3])
+        self.assertEqual(self.cdd[4], 9876)
 
     def test_one_in_disk_then_moves_to_memory_then_disk_again(self):
-        self.cdd[1] = None
-        self.cdd[2] = None
-        self.cdd[3] = None
+        self.cdd[1] = 1
+        self.cdd[2] = 2
+        self.cdd[3] = 3
 
-        self.cdd[1] = None
-        self.cdd[2] = None
-        self.cdd[3] = None
+        self.cdd[1] = 1
+        self.cdd[2] = 2
+        self.cdd[3] = 3
 
-        self.cdd[4] = None
+        self.cdd[4] = 4
 
         self.assertEqual(self.cdd._in_memory.keys(), [1, 2, 3])
         self.assertEqual(self.cdd._disk_dict.keys(), [4])
@@ -120,3 +129,6 @@ def test_one_in_disk_then_moves_to_memory_then_disk_again(self):
 
         self.assertEqual(self.cdd._in_memory.keys(), [1, 2, 3])
         self.assertEqual(self.cdd._disk_dict.keys(), [4])
+
+        self.assertEqual([(k, v) for (k, v) in self.cdd.iteritems()],
+                         [(1, 1), (2, 2), (3, 3), (4, 4)])