3232
3333import pytest
3434
35- pytestmark = [ pytest . mark . psycopg_pool , pytest . mark . integrationtest ]
35+ from elasticapm . conf . constants import SPAN
3636
37- pytest .importorskip ("psycopg" )
37+ psycopg = pytest .importorskip ("psycopg" )
3838pool_mod = pytest .importorskip ("psycopg_pool" )
3939
40+ pytestmark = [pytest .mark .psycopg_pool , pytest .mark .integrationtest ]
41+
42+ has_postgres_configured = "POSTGRES_DB" in os .environ
43+
44+
45+ def connect_kwargs ():
46+ return {
47+ "dbname" : os .environ .get ("POSTGRES_DB" , "elasticapm_test" ),
48+ "user" : os .environ .get ("POSTGRES_USER" , "postgres" ),
49+ "password" : os .environ .get ("POSTGRES_PASSWORD" , "postgres" ),
50+ "host" : os .environ .get ("POSTGRES_HOST" , None ),
51+ "port" : os .environ .get ("POSTGRES_PORT" , None ),
52+ }
53+
54+
55+ def make_conninfo ():
56+ kw = connect_kwargs ()
57+ host = kw ["host" ] or "localhost"
58+ port = kw ["port" ] or "5432"
59+ return f"postgresql://{ kw ['user' ]} :{ kw ['password' ]} @{ host } :{ port } /{ kw ['dbname' ]} "
60+
4061
62+ @pytest .mark .skipif (not has_postgres_configured , reason = "PostgreSQL not configured" )
4163def test_pool_generates_spans (instrument , elasticapm_client ):
4264 with pool_mod .ConnectionPool (
43- os .environ .get (
44- "PSYCOPG_TEST_DSN" ,
45- "postgresql://postgres:postgres@127.0.0.1:5432/postgres" ,
46- ),
65+ make_conninfo (),
4766 min_size = 1 ,
4867 max_size = 2 ,
4968 open = True ,
@@ -59,5 +78,8 @@ def test_pool_generates_spans(instrument, elasticapm_client):
5978 finally :
6079 elasticapm_client .end_transaction ("200" )
6180
62- spans = elasticapm_client .events ["span" ]
81+ spans = elasticapm_client .events [SPAN ]
82+ # Verify that connect span and query span are generated
83+ assert len (spans ) >= 2
84+ assert any (span .get ("action" ) == "connect" for span in spans )
6385 assert any (span .get ("context" , {}).get ("db" , {}).get ("type" ) == "sql" for span in spans )
0 commit comments