- modernize multiple warning assertions for PG reflect index test

2 files changed, 28 insertions, 24 deletions

diff --git a/test/dialect/test_postgresql.py b/test/dialect/test_postgresql.py
index 06a39ec26..10dedab46 100644
--- a/test/dialect/test_postgresql.py
+++ b/test/dialect/test_postgresql.py
@@ -1293,24 +1293,15 @@ class MiscTest(TestBase, AssertsExecutionResults, AssertsCompiledSQL):
         finally:
             testing.db.execute('drop table speedy_users')
 
-    @testing.emits_warning()
+    @testing.provide_metadata
     def test_index_reflection(self):
         """ Reflecting partial & expression-based indexes should warn
         """
 
-        import warnings
-
-        def capture_warnings(*args, **kw):
-            capture_warnings._orig_showwarning(*args, **kw)
-            capture_warnings.warnings.append(args)
-
-        capture_warnings._orig_showwarning = warnings.warn
-        capture_warnings.warnings = []
-        m1 = MetaData(testing.db)
-        t1 = Table('party', m1, Column('id', String(10),
+        t1 = Table('party', metadata, Column('id', String(10),
                    nullable=False), Column('name', String(20),
                    index=True), Column('aname', String(20)))
-        m1.create_all()
+        metadata.create_all()
         testing.db.execute("""
           create index idx1 on party ((id || name))
         """)
@@ -1321,17 +1312,10 @@ class MiscTest(TestBase, AssertsExecutionResults, AssertsCompiledSQL):
             create index idx3 on party using btree
                 (lower(name::text), lower(aname::text))
         """)
-        try:
+
+        def go():
             m2 = MetaData(testing.db)
-            warnings.warn = capture_warnings
             t2 = Table('party', m2, autoload=True)
-            wrn = capture_warnings.warnings
-            assert str(wrn[0][0]) \
-                == 'Skipped unsupported reflection of '\
-                'expression-based index idx1'
-            assert str(wrn[1][0]) \
-                == 'Predicate of partial index idx2 ignored during '\
-                'reflection'
             assert len(t2.indexes) == 2
 
             # Make sure indexes are in the order we expect them in
@@ -1344,9 +1328,17 @@ class MiscTest(TestBase, AssertsExecutionResults, AssertsCompiledSQL):
             assert r2.unique == False
             assert [t2.c.id] == r1.columns
             assert [t2.c.name] == r2.columns
-        finally:
-            warnings.warn = capture_warnings._orig_showwarning
-            m1.drop_all()
+
+        testing.assert_warnings(go,
+            [
+                'Skipped unsupported reflection of '
+                'expression-based index idx1',
+                'Predicate of partial index idx2 ignored during '
+                'reflection',
+                'Skipped unsupported reflection of '
+                'expression-based index idx3'
+            ])
+        
 
     @testing.fails_on('postgresql+pypostgresql',
                       'pypostgresql bombs on multiple calls')
diff --git a/test/lib/testing.py b/test/lib/testing.py
index 00b65c6e3..8f4838d0a 100644
--- a/test/lib/testing.py
+++ b/test/lib/testing.py
@@ -387,6 +387,18 @@ def emits_warning_on(db, *warnings):
         return function_named(maybe, fn.__name__)
     return decorate
 
+def assert_warnings(fn, warnings):
+    """Assert that each of the given warnings are emitted by fn."""
+    
+    orig_warn = util.warn
+    def capture_warnings(*args, **kw):
+        orig_warn(*args, **kw)
+        popwarn = warnings.pop(0)
+        eq_(args[0], popwarn)
+    util.warn = util.langhelpers.warn = capture_warnings
+    
+    return emits_warning()(fn)()
+    
 def uses_deprecated(*messages):
     """Mark a test as immune from fatal deprecation warnings.