- add support for assertion of warnings emitted

1 files changed, 117 insertions, 0 deletions

diff --git a/alembic/testing/assertions.py b/alembic/testing/assertions.py
index f5e992e..fd65fd6 100644
--- a/alembic/testing/assertions.py
+++ b/alembic/testing/assertions.py
@@ -1,7 +1,16 @@
+from __future__ import absolute_import
+
+
 import re
 from alembic import util
 from sqlalchemy.engine import default
 from alembic.compat import text_type, py3k
+import contextlib
+from sqlalchemy.util import decorator
+from sqlalchemy import exc as sa_exc
+import warnings
+from . import mock
+
 
 if not util.sqla_094:
     def eq_(a, b, msg=None):
@@ -82,3 +91,111 @@ def _get_dialect(name):
                 d.implicit_returning = True
             return d
 
+
+def expect_warnings(*messages, **kw):
+    """Context manager which expects one or more warnings.
+
+    With no arguments, squelches all SAWarnings emitted via
+    sqlalchemy.util.warn and sqlalchemy.util.warn_limited.   Otherwise
+    pass string expressions that will match selected warnings via regex;
+    all non-matching warnings are sent through.
+
+    The expect version **asserts** that the warnings were in fact seen.
+
+    Note that the test suite sets SAWarning warnings to raise exceptions.
+
+    """
+    return _expect_warnings(sa_exc.SAWarning, messages, **kw)
+
+
+@contextlib.contextmanager
+def expect_warnings_on(db, *messages, **kw):
+    """Context manager which expects one or more warnings on specific
+    dialects.
+
+    The expect version **asserts** that the warnings were in fact seen.
+
+    """
+    spec = db_spec(db)
+
+    if isinstance(db, util.string_types) and not spec(config._current):
+        yield
+    elif not _is_excluded(*db):
+        yield
+    else:
+        with expect_warnings(*messages, **kw):
+            yield
+
+
+def emits_warning(*messages):
+    """Decorator form of expect_warnings().
+
+    Note that emits_warning does **not** assert that the warnings
+    were in fact seen.
+
+    """
+
+    @decorator
+    def decorate(fn, *args, **kw):
+        with expect_warnings(assert_=False, *messages):
+            return fn(*args, **kw)
+
+    return decorate
+
+
+def emits_warning_on(db, *messages):
+    """Mark a test as emitting a warning on a specific dialect.
+
+    With no arguments, squelches all SAWarning failures.  Or pass one or more
+    strings; these will be matched to the root of the warning description by
+    warnings.filterwarnings().
+
+    Note that emits_warning_on does **not** assert that the warnings
+    were in fact seen.
+
+    """
+    @decorator
+    def decorate(fn, *args, **kw):
+        with expect_warnings_on(db, *messages):
+            return fn(*args, **kw)
+
+    return decorate
+
+
+@contextlib.contextmanager
+def _expect_warnings(exc_cls, messages, regex=True, assert_=True):
+
+    if regex:
+        filters = [re.compile(msg, re.I) for msg in messages]
+    else:
+        filters = messages
+
+    seen = set(filters)
+
+    real_warn = warnings.warn
+
+    def our_warn(msg, exception=None, *arg, **kw):
+        if exception and not issubclass(exception, exc_cls):
+            return real_warn(msg, exception, *arg, **kw)
+
+        if not filters:
+            return
+
+        for filter_ in filters:
+            if (regex and filter_.match(msg)) or \
+                    (not regex and filter_ == msg):
+                seen.discard(filter_)
+                break
+        else:
+            if exception is None:
+                real_warn(msg, *arg, **kw)
+            else:
+                real_warn(msg, exception, *arg, **kw)
+
+    with mock.patch("warnings.warn", our_warn):
+        yield
+
+    if assert_:
+        assert not seen, "Warnings were not seen: %s" % \
+            ", ".join("%r" % (s.pattern if regex else s) for s in seen)
+