Implemented support for "getitem" and "setitem" action kinds.

2 files changed, 53 insertions, 1 deletions

diff --git a/mocker.py b/mocker.py
index 7d0db79..058105d 100644
--- a/mocker.py
+++ b/mocker.py
@@ -754,6 +754,12 @@ class Mock(object):
     def __contains__(self, value):
         return self.__mocker_act__("contains", (value,))
 
+    def __getitem__(self, key):
+        return self.__mocker_act__("getitem", (key,))
+
+    def __setitem__(self, key, value):
+        return self.__mocker_act__("setitem", (key, value))
+
     # When adding a new action kind here, also add support for it on
     # Action.execute() and Path.__str__().
 
@@ -833,6 +839,10 @@ class Action(object):
                 result = object(*self.args, **self.kwargs)
             elif kind == "contains":
                 result = self.args[0] in object
+            elif kind == "getitem":
+                result = object[self.args[0]]
+            elif kind == "setitem":
+                result = object[self.args[0]] = self.args[1]
             else:
                 raise RuntimeError("Don't know how to execute %r kind." % kind)
         self._execute_cache[id(object)] = result
@@ -911,6 +921,11 @@ class Path(object):
                 result = "%s(%s)" % (result, ", ".join(args))
             elif action.kind == "contains":
                 result = "%r in %s" % (action.args[0], result)
+            elif action.kind == "getitem":
+                result = "%s[%r]" % (result, action.args[0])
+            elif action.kind == "setitem":
+                result = "%s[%r] = %r" % (result, action.args[0],
+                                          action.args[1])
             else:
                 raise RuntimeError("Don't know how to format kind %r" %
                                    action.kind)
diff --git a/test.py b/test.py
index b07d604..db71100 100755
--- a/test.py
+++ b/test.py
@@ -940,13 +940,26 @@ class ActionTest(unittest.TestCase):
         action = Action("call", (1,), {"b": 2})
         self.assertEquals(action.execute(obj), 3)
 
-    def test_contains(self):
+    def test_execute_contains(self):
         obj = set(["a"])
         action = Action("contains", ("a",), {})
         self.assertEquals(action.execute(obj), True)
         action = Action("contains", ("b",), {})
         self.assertEquals(action.execute(obj), False)
 
+    def test_execute_getitem(self):
+        obj = {"a": 1}
+        action = Action("getitem", ("a",), {})
+        self.assertEquals(action.execute(obj), 1)
+        action = Action("getitem", ("b",), {})
+        self.assertRaises(KeyError, action.execute, obj)
+
+    def test_execute_setitem(self):
+        obj = {}
+        action = Action("setitem", ("a", 1), {})
+        action.execute(obj)
+        self.assertEquals(obj, {"a": 1})
+
     def test_execute_caching(self):
         values = iter(range(10))
         obj = lambda: values.next()
@@ -1189,6 +1202,14 @@ class PathTest(unittest.TestCase):
         path = Path(self.mock, None, [Action("contains", ("value",), {})])
         self.assertEquals(str(path), "'value' in obj")
 
+    def test_str_getitem(self):
+        path = Path(self.mock, None, [Action("getitem", ("key",), {})])
+        self.assertEquals(str(path), "obj['key']")
+
+    def test_str_setitem(self):
+        path = Path(self.mock, None, [Action("setitem", ("key", "value"), {})])
+        self.assertEquals(str(path), "obj['key'] = 'value'")
+
     def test_str_getattr_call(self):
         path = Path(self.mock, None, [Action("getattr", ("x",), {}),
                                       Action("getattr", ("y",), {}),
@@ -1496,6 +1517,22 @@ class MockTest(unittest.TestCase):
         self.assertEquals(path, self.mock.__mocker_path__ + 
                                 Action("contains", ("value",), {}))
 
+    def test_getitem(self):
+        self.assertEquals(self.mock["key"], 42)
+        (path,) = self.paths
+        self.assertEquals(type(path), Path)
+        self.assertTrue(path.parent_path is self.mock.__mocker_path__)
+        self.assertEquals(path, self.mock.__mocker_path__ + 
+                                Action("getitem", ("key",), {}))
+
+    def test_setitem(self):
+        self.mock["key"] = "value"
+        (path,) = self.paths
+        self.assertEquals(type(path), Path)
+        self.assertTrue(path.parent_path is self.mock.__mocker_path__)
+        self.assertEquals(path, self.mock.__mocker_path__ + 
+                                Action("setitem", ("key", "value"), {}))
+
     def test_passthrough_on_unexpected(self):
         class StubMocker(object):
             def act(self, path):