Modify postgresql tests to include WITH query

tests/agent_features/test_span_events.py

@@ -141,7 +141,6 @@ def test_each_span_type(trace_type, args):
     )
     @background_task(name="test_each_span_type")
     def _test():
-
         transaction = current_transaction()
         transaction._sampled = True
 
@@ -307,7 +306,6 @@ def _test():
     }
 )
 def test_external_span_limits(kwarg_override, attr_override):
-
     exact_intrinsics = {
         "type": "Span",
         "sampled": True,
@@ -364,7 +362,6 @@ def _test():
     }
 )
 def test_datastore_span_limits(kwarg_override, attribute_override):
-
     exact_intrinsics = {
         "type": "Span",
         "sampled": True,
@@ -416,10 +413,6 @@ def _test():
 @pytest.mark.parametrize("span_events_enabled", (False, True))
 def test_collect_span_events_override(collect_span_events, span_events_enabled):
     spans_expected = collect_span_events and span_events_enabled
-    # if collect_span_events and span_events_enabled:
-    #     spans_expected = True
-    # else:
-    #     spans_expected = False
 
     span_count = 2 if spans_expected else 0
 
@@ -509,7 +502,6 @@ def __exit__(self, *args):
 )
 @pytest.mark.parametrize("exclude_attributes", (True, False))
 def test_span_event_user_attributes(trace_type, args, exclude_attributes):
-
     _settings = {
         "distributed_tracing.enabled": True,
         "span_events.enabled": True,
@@ -626,7 +618,6 @@ def _test():
     ),
 )
 def test_span_event_error_attributes_notice_error(trace_type, args):
-
     _settings = {
         "distributed_tracing.enabled": True,
         "span_events.enabled": True,
@@ -674,7 +665,6 @@ def _test():
     ),
 )
 def test_span_event_error_attributes_observed(trace_type, args):
-
     error = ValueError("whoops")
 
     exact_agents = {

tests/datastore_postgresql/conftest.py

@@ -12,21 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-
-from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture  # noqa: F401; pylint: disable=W0611
-
+from testing_support.fixtures import (  # noqa: F401; pylint: disable=W0611
+    collector_agent_registration_fixture,
+    collector_available_fixture,
+)
 
 _default_settings = {
-    'transaction_tracer.explain_threshold': 0.0,
-    'transaction_tracer.transaction_threshold': 0.0,
-    'transaction_tracer.stack_trace_threshold': 0.0,
-    'debug.log_data_collector_payloads': True,
-    'debug.record_transaction_failure': True,
-    'debug.log_explain_plan_queries': True
+    "transaction_tracer.explain_threshold": 0.0,
+    "transaction_tracer.transaction_threshold": 0.0,
+    "transaction_tracer.stack_trace_threshold": 0.0,
+    "debug.log_data_collector_payloads": True,
+    "debug.record_transaction_failure": True,
+    "debug.log_explain_plan_queries": True,
 }
 
 collector_agent_registration = collector_agent_registration_fixture(
-        app_name='Python Agent Test (datastore_postgresql)',
-        default_settings=_default_settings,
-        linked_applications=['Python Agent Test (datastore)'])
+    app_name="Python Agent Test (datastore_postgresql)",
+    default_settings=_default_settings,
+    linked_applications=["Python Agent Test (datastore)"],
+)

tests/datastore_postgresql/test_database.py

@@ -13,15 +13,13 @@
 # limitations under the License.
 
 import postgresql.driver.dbapi20
-
-
-from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
-
+from testing_support.db_settings import postgresql_settings
 from testing_support.validators.validate_database_trace_inputs import (
     validate_database_trace_inputs,
 )
-
-from testing_support.db_settings import postgresql_settings
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
 
 from newrelic.api.background_task import background_task
 
@@ -41,13 +39,14 @@
     ("Datastore/operation/Postgres/create", 1),
     ("Datastore/operation/Postgres/commit", 3),
     ("Datastore/operation/Postgres/rollback", 1),
+    ("Datastore/operation/Postgres/other", 1),
 ]
 
 _test_execute_via_cursor_rollup_metrics = [
-    ("Datastore/all", 13),
-    ("Datastore/allOther", 13),
-    ("Datastore/Postgres/all", 13),
-    ("Datastore/Postgres/allOther", 13),
+    ("Datastore/all", 14),
+    ("Datastore/allOther", 14),
+    ("Datastore/Postgres/all", 14),
+    ("Datastore/Postgres/allOther", 14),
     ("Datastore/operation/Postgres/select", 1),
     ("Datastore/statement/Postgres/%s/select" % DB_SETTINGS["table_name"], 1),
     ("Datastore/operation/Postgres/insert", 1),
@@ -63,6 +62,10 @@
     ("Datastore/operation/Postgres/call", 2),
     ("Datastore/operation/Postgres/commit", 3),
     ("Datastore/operation/Postgres/rollback", 1),
+    ("Datastore/operation/Postgres/other", 1),
+    ("Function/postgresql.driver.dbapi20:connect", 1),
+    ("Function/postgresql.driver.dbapi20:Connection.__enter__", 1),
+    ("Function/postgresql.driver.dbapi20:Connection.__exit__", 1),
 ]
 
 
@@ -82,30 +85,27 @@ def test_execute_via_cursor():
         host=DB_SETTINGS["host"],
         port=DB_SETTINGS["port"],
     ) as connection:
-
         cursor = connection.cursor()
 
         cursor.execute("""drop table if exists %s""" % DB_SETTINGS["table_name"])
 
-        cursor.execute(
-            """create table %s """ % DB_SETTINGS["table_name"]
-            + """(a integer, b real, c text)"""
-        )
+        cursor.execute("""create table %s """ % DB_SETTINGS["table_name"] + """(a integer, b real, c text)""")
 
         cursor.executemany(
-            """insert into %s """ % DB_SETTINGS["table_name"]
-            + """values (%s, %s, %s)""",
+            """insert into %s """ % DB_SETTINGS["table_name"] + """values (%s, %s, %s)""",
             [(1, 1.0, "1.0"), (2, 2.2, "2.2"), (3, 3.3, "3.3")],
         )
 
         cursor.execute("""select * from %s""" % DB_SETTINGS["table_name"])
 
-        for row in cursor:
-            pass
+        cursor.execute(
+            """with temporaryTable (averageValue) as (select avg(b) from %s) """ % DB_SETTINGS["table_name"]
+            + """select * from %s,temporaryTable """ % DB_SETTINGS["table_name"]
+            + """where %s.b > temporaryTable.averageValue""" % DB_SETTINGS["table_name"]
+        )
 
         cursor.execute(
-            """update %s """ % DB_SETTINGS["table_name"]
-            + """set a=%s, b=%s, c=%s where a=%s""",
+            """update %s """ % DB_SETTINGS["table_name"] + """set a=%s, b=%s, c=%s where a=%s""",
             (4, 4.0, "4.0", 1),
         )
 
@@ -152,7 +152,6 @@ def test_rollback_on_exception():
             host=DB_SETTINGS["host"],
             port=DB_SETTINGS["port"],
         ):
-
             raise RuntimeError("error")
 
     except RuntimeError: