Adapt pg8000 tests to resemble the ones for psycopg2

aws · Jan 4, 2019 · 6b31f67 · 6b31f67
1 parent 5694ce8
commit 6b31f67
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/ext/pg8000/test_pg8000.py b/tests/ext/pg8000/test_pg8000.py
@@ -6,7 +6,8 @@
 from aws_xray_sdk.core import patch
 from aws_xray_sdk.core import xray_recorder
 from aws_xray_sdk.core.context import Context
-from aws_xray_sdk.ext.pg8000 import unpatch
+
+patch(('pg8000',))
 
 
 @pytest.fixture(autouse=True)
@@ -16,13 +17,11 @@ def construct_ctx():
     so that later subsegment can be attached. After each test run
     it cleans up context storage again.
     """
-    patch(('pg8000',))
     xray_recorder.configure(service='test', sampling=False, context=Context())
     xray_recorder.clear_trace_entities()
     xray_recorder.begin_segment('name')
     yield
     xray_recorder.clear_trace_entities()
-    unpatch()
 
 
 def test_execute_dsn_kwargs():