Enhance MockContext a whole bunch

- repeat kwarg
- boolean, string result value shorthands
- supports tuple values, not just lists
- regex result-dict keys
- other minor cleanup, fixes

Author: bitprophet

invoke/context.py

@@ -1,11 +1,12 @@
 import os
 import re
 from contextlib import contextmanager
+from itertools import cycle
 
 try:
-    from invoke.vendor.six import raise_from, iteritems
+    from invoke.vendor.six import raise_from, iteritems, string_types
 except ImportError:
-    from six import raise_from, iteritems
+    from six import raise_from, iteritems, string_types
 
 from .config import Config, DataProxy
 from .exceptions import Failure, AuthFailure, ResponseNotAccepted
@@ -404,26 +405,55 @@ def __init__(self, config=None, **kwargs):
             A Configuration object to use. Identical in behavior to `.Context`.
 
         :param run:
-            A data structure of `Results <.Result>`, to return from calls to
-            the instantiated object's `~.Context.run` method (instead of
-            actually executing the requested shell command).
+            A data structure indicating what `.Result` objects to return from
+            calls to the instantiated object's `~.Context.run` method (instead
+            of actually executing the requested shell command).
 
             Specifically, this kwarg accepts:
 
-            - A single `.Result` object, which will be returned once.
-            - An iterable of `Results <.Result>`, which will be returned on
-              each subsequent call to ``.run``.
-            - A map of command strings to either of the above, allowing
-              specific call-and-response semantics instead of assuming a call
-              order.
+            - A single `.Result` object.
+
+                - Remember that class's first positional argument is its stdout
+                  - it can thus be handy to hand in expressions like
+                    ``Result("command output here!")``.)
+
+            - A boolean; if True, yields a `.Result` whose ``exited`` is ``0``,
+              and if False, ``1``.
+            - An iterable of the above values, which will be returned on each
+              subsequent call to ``.run`` (the first item on the first call,
+              the second on the second call, etc).
+            - A dict mapping command strings or compiled regexen to the above
+              values (including an iterable), allowing specific
+              call-and-response semantics instead of assuming a call order.
 
         :param sudo:
             Identical to ``run``, but whose values are yielded from calls to
             `~.Context.sudo`.
 
+        :param bool repeat:
+            A flag determining whether results yielded by this class' methods
+            repeat or are consumed.
+
+            For example, when a single result is indicated, it will normally
+            only be returned once, causing `NotImplementedError` afterwards.
+            But when ``repeat=True`` is given, that result is returned on
+            every call, forever.
+
+            Similarly, iterable results are normally exhausted once, but when
+            this setting is enabled, they are wrapped in `itertools.cycle`.
+
+            Default: ``False`` (for backwards compatibility reasons).
+
         :raises:
             ``TypeError``, if the values given to ``run`` or other kwargs
-            aren't individual `.Result` objects or iterables.
+            aren't of the expected types.
+
+        .. versionchanged:: 1.5
+            Added support for boolean and string result values.
+        .. versionchanged:: 1.5
+            Added support for regex dict keys.
+        .. versionchanged:: 1.5
+            Added the ``repeat`` keyword argument.
         """
         # Figure out if we can support Mock in the current environment
         Mock = None
@@ -434,52 +464,77 @@ def __init__(self, config=None, **kwargs):
                 from unittest.mock import Mock
             except ImportError:
                 pass
-        # TODO: would be nice to allow regexen instead of exact string matches
+        # Set up like any other Context would, with the config
         super(MockContext, self).__init__(config)
+        # Pull out behavioral kwargs
+        self._set("__repeat", kwargs.pop("repeat", False))
+        # The rest must be things like run/sudo - mock Context method info
         for method, results in iteritems(kwargs):
-            # Special convenience case: individual Result -> one-item list
-            if (
-                not hasattr(results, "__iter__")
-                and not isinstance(results, Result)
-                # No need for explicit dict test; they have __iter__
-            ):
+            # For each possible value type, normalize to iterable of Result
+            # objects (possibly repeating).
+            singletons = tuple([Result, bool] + list(string_types))
+            if isinstance(results, dict):
+                for key, value in iteritems(results):
+                    results[key] = self._normalize(value)
+            elif isinstance(results, singletons) or hasattr(results, "__iter__"):
+                results = self._normalize(results)
+            # Unknown input value: cry
+            else:
                 err = "Not sure how to yield results from a {!r}"
                 raise TypeError(err.format(type(results)))
-            # Set the return values
+            # Save results for use by the method
             self._set("__{}".format(method), results)
             # Wrap the method in a Mock, if applicable
             if Mock is not None:
                 self._set(method, Mock(wraps=getattr(self, method)))
 
+    def _normalize(self, value):
+        # First turn everything into an iterable
+        if not hasattr(value, "__iter__") or isinstance(value, string_types):
+            value = [value]
+        # Then turn everything within into a Result
+        results = []
+        for obj in value:
+            if isinstance(obj, bool):
+                obj = Result(exited=0 if obj else 1)
+            elif isinstance(obj, string_types):
+                obj = Result(obj)
+            results.append(obj)
+        # Finally, turn that iterable into an iteratOR, depending on repeat
+        return cycle(results) if getattr(self, "__repeat") else iter(results)
+
     # TODO: _maybe_ make this more metaprogrammy/flexible (using __call__ etc)?
     # Pretty worried it'd cause more hard-to-debug issues than it's presently
     # worth. Maybe in situations where Context grows a _lot_ of methods (e.g.
     # in Fabric 2; though Fabric could do its own sub-subclass in that case...)
 
     def _yield_result(self, attname, command):
-        # NOTE: originally had this with a bunch of explicit
-        # NotImplementedErrors, but it doubled method size, and chance of
-        # unexpected index/etc errors seems low here.
         try:
-            # Obtain result if possible
-            value = getattr(self, attname)
-            # TODO: thought there's a 'better' 2x3 DictType or w/e, but can't
-            # find one offhand
-            if isinstance(value, dict):
-                if hasattr(value[command], "__iter__"):
-                    result = value[command].pop(0)
-                elif isinstance(value[command], Result):
-                    result = value.pop(command)
-            elif hasattr(value, "__iter__"):
-                result = value.pop(0)
-            elif isinstance(value, Result):
-                result = value
-                delattr(self, attname)
-            # Populate command string unless explicitly given
+            obj = getattr(self, attname)
+            # Dicts need to try direct lookup or regex matching
+            if isinstance(obj, dict):
+                try:
+                    obj = obj[command]
+                except KeyError:
+                    # TODO: could optimize by skipping this if not any regex
+                    # objects in keys()?
+                    for key, value in iteritems(obj):
+                        if hasattr(key, "match") and key.match(command):
+                            obj = value
+                            break
+                    else:
+                        # Nope, nothing did match.
+                        raise KeyError
+            # Here, the value was either never a dict or has been extracted
+            # from one, so we can assume it's an iterable of Result objects due
+            # to work done by __init__.
+            result = next(obj)
+            # Populate Result's command string with what matched unless
+            # explicitly given
             if not result.command:
                 result.command = command
             return result
-        except (AttributeError, IndexError, KeyError):
+        except (AttributeError, IndexError, KeyError, StopIteration):
             raise_from(NotImplementedError, None)
 
     def run(self, command, *args, **kwargs):
@@ -531,4 +586,4 @@ def set_result_for(self, attname, command, result):
         if not isinstance(value, dict):
             raise heck
         # OK, we're good to modify, so do so.
-        value[command] = result
+        value[command] = self._normalize(result)

sites/docs/concepts/testing.rst

@@ -97,12 +97,120 @@ possible commands, using a dict value::
     def test_homebrew_gsed():
         expected_sed = "gsed -e s/foo/bar/g file.txt"
         c = MockContext(run={
-            "which gsed": Result(),
+            "which gsed": Result(exited=0),
             expected_sed: Result(),
         })
         replace(c, 'file.txt', 'foo', 'bar')
         c.run.assert_called_with(expected_sed)
 
+Boolean mock results
+--------------------
+
+You may have noticed the above example uses a handful of 'empty' `.Result`
+objects; these stand in for "succeeded, but otherwise had no useful attributes"
+command executions (as `.Result` defaults to an exit code of ``0`` and empty
+strings for stdout/stderr).
+
+This is relatively common - think "interrogative" commands where the caller
+only cares for a boolean result, or times when a command is called purely for
+its side effects. To support this, there's a shorthand in `.MockContext`:
+passing ``True`` or ``False`` to stand in for otherwise blank Results with exit
+codes of ``0`` or ``1`` respectively.
+
+The example tests then look like this::
+
+    from invoke import MockContext, Result
+    from mytasks import replace
+
+    def test_regular_sed():
+        expected_sed = "sed -e s/foo/bar/g file.txt"
+        c = MockContext(run={
+            "which gsed": False,
+            expected_sed: True,
+        })
+        replace(c, 'file.txt', 'foo', 'bar')
+        c.run.assert_called_with(expected_sed)
+
+    def test_homebrew_gsed():
+        expected_sed = "gsed -e s/foo/bar/g file.txt"
+        c = MockContext(run={
+            "which gsed": True,
+            expected_sed: True,
+        })
+        replace(c, 'file.txt', 'foo', 'bar')
+        c.run.assert_called_with(expected_sed)
+
+String mock results
+-------------------
+
+Another convenient shorthand is using string values, which are interpreted to
+be the stdout of the resulting `.Result`. This only really saves you from
+writing out the class itself (since ``stdout`` is the first positional arg of
+`.Result`!) but "command X results in stdout Y" is a common enough use case
+that we implemented it anyway.
+
+By example, let's modify an earlier example where we cared about stdout::
+
+    from invoke import MockContext
+    from mytasks import get_platform
+
+    def test_get_platform_on_mac():
+        c = MockContext(run="Darwin\n")
+        assert "Apple" in get_platform(c)
+
+    def test_get_platform_on_linux():
+        c = MockContext(run="Linux\n")
+        assert "desktop" in get_platform(c)
+
+As with everything else in this document, this tactic can be applied to
+iterators or mappings as well as individual values.
+
+Regular expression command matching
+-----------------------------------
+
+The dict form of `.MockContext` kwarg can accept regular expression objects as
+keys, in addition to strings; ideal for situations where you either don't know
+the exact command being invoked, or simply don't need or want to write out the
+entire thing.
+
+Imagine you're writing a function to run package management commands on a few
+different Linux distros and you're trying to test its error handling. You might
+want to set up a context that pretends any arbitrary ``apt`` or ``yum`` command
+fails, and ensure the function returns stderr when it encounters a problem::
+
+    import re
+    from invoke import MockContext
+    from mypackage.tasks import install
+
+    package_manager = re.compile(r"^(apt(-get)?|yum) .*")
+
+    def test_package_success_returns_True():
+        c = MockContext(run={package_manager: True})
+        assert install(c, package="somepackage") is True
+
+    def test_package_explosions_return_stderr():
+        c = MockContext(run={
+            package_manager: Result(stderr="oh no!", exited=1),
+        })
+        assert install(c, package="otherpackage") == "oh no!"
+
+A bit contrived - there are a bunch of other ways to organize this exact test
+code so you don't truly need the regex - but hopefully it's clear that when you
+*do* need this flexibility, this is how you could go about it.
+
+Repeated results
+----------------
+
+By default, the values in these mock structures are consumed, causing
+`.MockContext` to raise ``NotImplementedError`` afterwards (as it does for any
+unexpected command executions). This was designed with the assumption that
+most code under test will run a given command once.
+
+If your situation doesn't match this, give ``repeat=True`` to the constructor,
+and you'll see values repeat indefinitely instead (or in cycles, for
+iterables).
+
+
 Expect `Results <.Result>`
 ==========================
 

sites/www/changelog.rst

@@ -2,6 +2,26 @@
 Changelog
 =========
 
+- :feature:`-` `~invoke.context.MockContext` now accepts a few quality-of-life
+  shortcuts as keys and values in its ``run``/``sudo`` arguments:
+
+    - Keys may be compiled regular expression objects, as well as strings, and
+      will match any calls whose commands match the regex.
+    - Values may be ``True`` or ``False`` as shorthand for otherwise empty
+      `~invoke.runners.Result` objects with exit codes of ``0`` or ``1``
+      respectively.
+    - Values may also be strings, as shorthand for otherwise empty
+      `~invoke.runners.Result` objects with those strings given as the
+      ``stdout`` argument.
+
+- :feature:`441` Add a new ``repeat`` kwarg to `~invoke.context.MockContext`
+  which, when True (default: False) causes stored results for its methods to be
+  yielded repeatedly instead of consumed. Feature request courtesy of
+  ``@SwampFalc``.
+- :bug:`- major` Immutable iterable result values handed to
+  `~invoke.context.MockContext` would yield errors (due to the use of
+  ``pop()``). The offending logic has been retooled to be more iterator-focused
+  and now works for tuples and etc.
 - :support:`-` Update the `testing documentation </concepts/testing>` a bit:
   cleaned up existing examples and added new sections for the other updates in
   the 1.5 release.