Commits

Konstantine Rybnikov committed 30d7ef6

add sequence_side_effect

  • Participants
  • Parent commits 70ad39e

Comments (0)

Files changed (3)

 README.rst
 mockstar.py
 requirements.txt
-setup.cfg
 setup.py
     """
     :class:`~MagicMock` with shortcuts. I just couldn't stand the
     ``.return_value.foo.return_value.bar.return_value`` thing.
+
+    Now you can do:
+
+    .. code-block:: python
+
+        from mockstar import M
+        m = M()
+        m.rv = 10
+        assert m() == 10
     """
     def __init__(self, *args, **kw):
         super(M, self).__init__(*args, **kw)
 
     #: shortcut for ``.return_value``
     rv = RVDescriptor()
+
+
+def sequence_side_effect(*args):
+    """
+    Generates function that returns i'th arg on i'th call.
+
+    Use like this:
+
+    .. code-block:: python
+
+        from mockstar import sequence_side_effect
+        from mockstar import M
+
+        m = M()
+        m.side_effect = sequence_side_effect(1, 2, 3)
+    """
+    seq = list(args)
+
+    def rv_fun(*args, **kw):
+        return seq.pop(0)
+    return rv_fun

File mockstar_test.py

 from mockstar import DotDict
 from mockstar import prefixed_p
 from mockstar import M
+from mockstar import sequence_side_effect
 
 
 def side_effect_one():
         self.assertIsInstance(se.side_effect_two, MagicMock)
 
 
+class TestSequenceSideEffect(TestCase):
+    def test_should_get_sequence(self):
+        m = M(side_effect=sequence_side_effect(1, 2, 3))
+        self.assertEquals(m(), 1)
+        self.assertEquals(m(), 2)
+        self.assertEquals(m(), 3)
+
+
 if __name__ == '__main__':
     unittest.main()