Skip to content

Commit 62c2bf0

Browse files
kgala2kgala2
committed
test: remove psycopg resource leak test and add remaining system tests
1 parent 177a776 commit 62c2bf0

1 file changed

Lines changed: 195 additions & 127 deletions

File tree

‎tests/system/test_psycopg_connection.py‎

Lines changed: 195 additions & 127 deletions
Original file line numberDiff line numberDiff line change
@@ -1,154 +1,222 @@
1-
# Copyright 2026 Google LLC
2-
#
3-
# Licensed under the Apache License, Version 2.0 (the "License");
4-
# you may not use this file except in compliance with the License.
5-
# You may obtain a copy of the License at
6-
#
7-
# http://www.apache.org/licenses/LICENSE-2.0
8-
#
9-
# Unless required by applicable law or agreed to in writing, software
10-
# distributed under the License is distributed on an "AS IS" BASIS,
11-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12-
# See the License for the specific language governing permissions and
13-
# limitations under the License.
1+
"""
2+
Copyright 2026 Google LLC
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
https://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
"""
1416

1517
import asyncio
18+
from datetime import datetime
1619
import os
17-
import time
18-
19-
try:
20-
import psutil
21-
except ImportError:
22-
psutil = None
20+
from typing import Union
2321

2422
import pytest
23+
import sqlalchemy
2524

2625
from google.cloud.sql.connector import Connector
27-
28-
# These will be set from environment variables or default to our test instance
29-
INSTANCE_CONNECTION_NAME = os.getenv(
30-
"DB_CONNECTION_NAME", "galakp-playground:us-east7:pg-us-east7-psycopg"
31-
)
32-
DB_USER = os.getenv("DB_USER", "postgres")
33-
DB_PASSWORD = os.getenv("DB_PASSWORD", "SuperPass123!")
34-
DB_NAME = os.getenv("DB_NAME", "postgres")
35-
36-
37-
@pytest.mark.skipif(psutil is None, reason="psutil package is not installed")
38-
def test_system_psycopg_resource_leak() -> None:
39-
"""Benchmark test to verify no resource leaks (threads, FDs, memory) between iteration 20 and 100."""
40-
print("\nStarting resource leak benchmark...")
41-
42-
process = psutil.Process(os.getpid())
43-
44-
def get_metrics():
45-
return {
46-
"threads": process.num_threads(),
47-
"fds": process.num_fds(),
48-
"rss_mb": process.memory_info().rss / (1024 * 1024),
49-
}
50-
51-
warmup_iterations = 20
52-
total_iterations = 100
53-
54-
baseline_metrics = None
55-
active_final_metrics = None
56-
57-
with Connector() as connector:
58-
for i in range(1, total_iterations + 1):
59-
conn = connector.connect(
60-
INSTANCE_CONNECTION_NAME,
61-
"psycopg",
62-
user=DB_USER,
63-
password=DB_PASSWORD,
64-
db=DB_NAME,
65-
)
66-
cursor = conn.cursor()
67-
cursor.execute("SELECT 1;")
68-
cursor.fetchone()
69-
cursor.close()
70-
conn.close()
71-
72-
if i == warmup_iterations:
73-
time.sleep(0.5)
74-
baseline_metrics = get_metrics()
75-
print(f"Baseline Metrics (Iteration {i}): {baseline_metrics}")
76-
77-
if i == total_iterations:
78-
time.sleep(0.5)
79-
active_final_metrics = get_metrics()
80-
print(f"Active Final Metrics (Iteration {i}): {active_final_metrics}")
81-
82-
if i % 10 == 0 and warmup_iterations < i < total_iterations:
83-
current_metrics = get_metrics()
84-
print(f"Iteration {i:3d}/{total_iterations}: {current_metrics}")
85-
86-
# Post-close metrics
87-
time.sleep(1)
88-
post_close_metrics = get_metrics()
89-
print(f"Post-Close Metrics: {post_close_metrics}")
90-
91-
assert baseline_metrics is not None
92-
assert active_final_metrics is not None
93-
94-
# Assertions: compare Active Final (100) vs Baseline (20)
95-
# Threads should not grow
96-
assert active_final_metrics["threads"] <= baseline_metrics["threads"] + 1, f"Thread leak: {baseline_metrics} -> {active_final_metrics}"
97-
# FDs should not grow
98-
assert active_final_metrics["fds"] <= baseline_metrics["fds"] + 1, f"FD leak: {baseline_metrics} -> {active_final_metrics}"
99-
# Memory growth should be minimal (allow < 5MB growth for minor fragmentation)
100-
assert active_final_metrics["rss_mb"] <= baseline_metrics["rss_mb"] + 5, f"Memory leak: {baseline_metrics} -> {active_final_metrics}"
101-
102-
print("Resource leak benchmark passed successfully.")
103-
104-
105-
def test_system_psycopg_basic() -> None:
106-
"""Basic system test to verify connection and query."""
107-
print(f"\nConnecting to {INSTANCE_CONNECTION_NAME}...")
108-
with Connector() as connector:
109-
conn = connector.connect(
110-
INSTANCE_CONNECTION_NAME,
26+
from google.cloud.sql.connector import DefaultResolver
27+
from google.cloud.sql.connector import DnsResolver
28+
29+
30+
def create_sqlalchemy_engine(
31+
instance_connection_name: str,
32+
user: str,
33+
password: str,
34+
db: str,
35+
ip_type: str = "public",
36+
refresh_strategy: str = "background",
37+
resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
38+
) -> tuple[sqlalchemy.engine.Engine, Connector]:
39+
"""Creates a connection pool for a Cloud SQL instance and returns the pool
40+
and the connector.
41+
"""
42+
connector = Connector(refresh_strategy=refresh_strategy, resolver=resolver)
43+
44+
# create SQLAlchemy connection pool
45+
engine = sqlalchemy.create_engine(
46+
"postgresql+psycopg://",
47+
creator=lambda: connector.connect(
48+
instance_connection_name,
11149
"psycopg",
112-
user=DB_USER,
113-
password=DB_PASSWORD,
114-
db=DB_NAME,
115-
)
116-
117-
cursor = conn.cursor()
118-
cursor.execute("SELECT version();")
119-
result = cursor.fetchone()
120-
print(f"Database version: {result[0]}")
121-
assert result is not None
122-
cursor.close()
123-
conn.close()
124-
print("Connection closed successfully.")
125-
126-
127-
50+
user=user,
51+
password=password,
52+
db=db,
53+
ip_type=ip_type,
54+
),
55+
)
56+
return engine, connector
57+
58+
59+
# Fallback to playground values if env vars are missing
60+
def get_env(key: str, default: str = "") -> str:
61+
# Map standard env vars to our playground values as defaults
62+
defaults = {
63+
"POSTGRES_CONNECTION_NAME": "galakp-playground:us-east7:pg-us-east7-psycopg",
64+
"POSTGRES_USER": "postgres",
65+
"POSTGRES_PASS": "SuperPass123!",
66+
"POSTGRES_DB": "postgres",
67+
}
68+
return os.getenv(key, defaults.get(key, default))
69+
70+
71+
def test_psycopg_connection() -> None:
72+
"""Basic test to get time from database using psycopg."""
73+
inst_conn_name = get_env("POSTGRES_CONNECTION_NAME")
74+
user = get_env("POSTGRES_USER")
75+
password = get_env("POSTGRES_PASS")
76+
db = get_env("POSTGRES_DB")
77+
ip_type = os.getenv("IP_TYPE", "public")
78+
79+
engine, connector = create_sqlalchemy_engine(
80+
inst_conn_name, user, password, db, ip_type
81+
)
82+
with engine.connect() as conn:
83+
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
84+
conn.commit()
85+
curr_time = time[0]
86+
assert type(curr_time) is datetime
87+
connector.close()
88+
89+
90+
def test_lazy_psycopg_connection() -> None:
91+
"""Basic test to get time from database using psycopg and lazy refresh."""
92+
inst_conn_name = get_env("POSTGRES_CONNECTION_NAME")
93+
user = get_env("POSTGRES_USER")
94+
password = get_env("POSTGRES_PASS")
95+
db = get_env("POSTGRES_DB")
96+
ip_type = os.getenv("IP_TYPE", "public")
97+
98+
engine, connector = create_sqlalchemy_engine(
99+
inst_conn_name, user, password, db, ip_type, "lazy"
100+
)
101+
with engine.connect() as conn:
102+
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
103+
conn.commit()
104+
curr_time = time[0]
105+
assert type(curr_time) is datetime
106+
connector.close()
107+
108+
109+
def test_CAS_psycopg_connection() -> None:
110+
"""Basic test to get time from database using CAS."""
111+
inst_conn_name = os.environ.get("POSTGRES_CAS_CONNECTION_NAME")
112+
user = get_env("POSTGRES_USER")
113+
password = os.environ.get("POSTGRES_CAS_PASS")
114+
db = get_env("POSTGRES_DB")
115+
ip_type = os.getenv("IP_TYPE", "public")
116+
117+
if not inst_conn_name or not password:
118+
pytest.skip("POSTGRES_CAS_CONNECTION_NAME or POSTGRES_CAS_PASS not set")
119+
120+
engine, connector = create_sqlalchemy_engine(
121+
inst_conn_name, user, password, db, ip_type
122+
)
123+
with engine.connect() as conn:
124+
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
125+
conn.commit()
126+
curr_time = time[0]
127+
assert type(curr_time) is datetime
128+
connector.close()
129+
130+
131+
def test_customer_managed_CAS_psycopg_connection() -> None:
132+
"""Basic test to get time from database using Customer Managed CAS."""
133+
inst_conn_name = os.environ.get("POSTGRES_CUSTOMER_CAS_CONNECTION_NAME")
134+
user = get_env("POSTGRES_USER")
135+
password = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS")
136+
db = get_env("POSTGRES_DB")
137+
ip_type = os.getenv("IP_TYPE", "public")
138+
139+
if not inst_conn_name or not password:
140+
pytest.skip("POSTGRES_CUSTOMER_CAS_CONNECTION_NAME or POSTGRES_CUSTOMER_CAS_PASS not set")
141+
142+
engine, connector = create_sqlalchemy_engine(
143+
inst_conn_name, user, password, db, ip_type
144+
)
145+
with engine.connect() as conn:
146+
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
147+
conn.commit()
148+
curr_time = time[0]
149+
assert type(curr_time) is datetime
150+
connector.close()
151+
152+
153+
def test_custom_SAN_with_dns_psycopg_connection() -> None:
154+
"""Basic test to get time from database using Custom SAN with DNS."""
155+
inst_conn_name = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME")
156+
user = get_env("POSTGRES_USER")
157+
password = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS")
158+
db = get_env("POSTGRES_DB")
159+
ip_type = os.getenv("IP_TYPE", "public")
160+
161+
if not inst_conn_name or not password:
162+
pytest.skip("POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME or POSTGRES_CUSTOMER_CAS_PASS not set")
163+
164+
engine, connector = create_sqlalchemy_engine(
165+
inst_conn_name, user, password, db, ip_type, resolver=DnsResolver
166+
)
167+
with engine.connect() as conn:
168+
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
169+
conn.commit()
170+
curr_time = time[0]
171+
assert type(curr_time) is datetime
172+
connector.close()
173+
174+
175+
def test_MCP_psycopg_connection() -> None:
176+
"""Basic test to get time from database using MCP enabled instance."""
177+
inst_conn_name = os.environ.get("POSTGRES_MCP_CONNECTION_NAME")
178+
user = get_env("POSTGRES_USER")
179+
password = os.environ.get("POSTGRES_MCP_PASS")
180+
db = get_env("POSTGRES_DB")
181+
ip_type = os.getenv("IP_TYPE", "public")
182+
183+
if not inst_conn_name or not password:
184+
pytest.skip("POSTGRES_MCP_CONNECTION_NAME or POSTGRES_MCP_PASS not set")
185+
186+
engine, connector = create_sqlalchemy_engine(
187+
inst_conn_name, user, password, db, ip_type
188+
)
189+
with engine.connect() as conn:
190+
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
191+
conn.commit()
192+
curr_time = time[0]
193+
assert type(curr_time) is datetime
194+
connector.close()
128195

129196

130197
def test_system_psycopg_to_thread() -> None:
131198
"""Verify that running sync connect in asyncio.to_thread works."""
132-
print(f"\nConnecting via asyncio.to_thread to {INSTANCE_CONNECTION_NAME}...")
199+
inst_conn_name = get_env("POSTGRES_CONNECTION_NAME")
200+
user = get_env("POSTGRES_USER")
201+
password = get_env("POSTGRES_PASS")
202+
db = get_env("POSTGRES_DB")
133203

134204
async def run_connect():
135205
with Connector() as connector:
136206
# Run the blocking connector.connect in a thread
137207
conn = await asyncio.to_thread(
138208
connector.connect,
139-
INSTANCE_CONNECTION_NAME,
209+
inst_conn_name,
140210
"psycopg",
141-
user=DB_USER,
142-
password=DB_PASSWORD,
143-
db=DB_NAME,
211+
user=user,
212+
password=password,
213+
db=db,
144214
)
145215
cursor = conn.cursor()
146-
cursor.execute("SELECT version();")
216+
cursor.execute("SELECT NOW();")
147217
result = cursor.fetchone()
148-
print(f"Database version (to_thread): {result[0]}")
149218
assert result is not None
150219
cursor.close()
151220
conn.close()
152221

153222
asyncio.run(run_connect())
154-
print("to_thread connection closed successfully.")

0 commit comments

Comments
 (0)