Skip to content

Commit 6f06ebf

Browse files
kgala2kgala2
committed
chore: update ip_type parameter in system tests
1 parent eba2228 commit 6f06ebf

8 files changed

+75
-44
lines changed

‎tests/system/test_asyncpg_connection.py‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ async def create_sqlalchemy_engine(
3232
user: str,
3333
password: str,
3434
db: str,
35+
ip_type: str,
3536
refresh_strategy: str = "background",
3637
resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
3738
) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, Connector]:
@@ -63,6 +64,8 @@ async def create_sqlalchemy_engine(
6364
The database user's password, e.g., secret-password
6465
db (str):
6566
The name of the database, e.g., mydb
67+
ip_type (str):
68+
The IP type of the Cloud SQL instance.
6669
refresh_strategy (Optional[str]):
6770
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6871
or "background". For serverless environments use "lazy" to avoid
@@ -87,9 +90,7 @@ async def create_sqlalchemy_engine(
8790
user=user,
8891
password=password,
8992
db=db,
90-
ip_type=os.environ.get(
91-
"IP_TYPE", "public"
92-
), # can also be "private" or "psc"
93+
ip_type=ip_type,
9394
),
9495
execution_options={"isolation_level": "AUTOCOMMIT"},
9596
)
@@ -101,6 +102,7 @@ async def create_asyncpg_pool(
101102
user: str,
102103
password: str,
103104
db: str,
105+
ip_type: str,
104106
refresh_strategy: str = "background",
105107
) -> tuple[asyncpg.Pool, Connector]:
106108
"""Creates a native asyncpg connection pool for a Cloud SQL instance and
@@ -130,6 +132,8 @@ async def create_asyncpg_pool(
130132
The database user's password, e.g., secret-password
131133
db (str):
132134
The name of the database, e.g., mydb
135+
ip_type (str):
136+
The IP type of the Cloud SQL instance.
133137
refresh_strategy (Optional[str]):
134138
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
135139
or "background". For serverless environments use "lazy" to avoid
@@ -147,9 +151,7 @@ async def getconn(
147151
user=user,
148152
password=password,
149153
db=db,
150-
ip_type=os.environ.get(
151-
"IP_TYPE", "public"
152-
), # can also be "private" or "psc",
154+
ip_type=ip_type
153155
**kwargs,
154156
)
155157
return conn
@@ -165,8 +167,9 @@ async def test_sqlalchemy_connection_with_asyncpg() -> None:
165167
user = os.environ["POSTGRES_USER"]
166168
password = os.environ["POSTGRES_PASS"]
167169
db = os.environ["POSTGRES_DB"]
170+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
168171

169-
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, password, db)
172+
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, password, db, ip_type)
170173

171174
async with pool.connect() as conn:
172175
res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone()
@@ -181,9 +184,10 @@ async def test_lazy_sqlalchemy_connection_with_asyncpg() -> None:
181184
user = os.environ["POSTGRES_USER"]
182185
password = os.environ["POSTGRES_PASS"]
183186
db = os.environ["POSTGRES_DB"]
187+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
184188

185189
pool, connector = await create_sqlalchemy_engine(
186-
inst_conn_name, user, password, db, "lazy"
190+
inst_conn_name, user, password, db, ip_type, "lazy"
187191
)
188192

189193
async with pool.connect() as conn:
@@ -199,9 +203,10 @@ async def test_custom_SAN_with_dns_sqlalchemy_connection_with_asyncpg() -> None:
199203
user = os.environ["POSTGRES_USER"]
200204
password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
201205
db = os.environ["POSTGRES_DB"]
206+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
202207

203208
pool, connector = await create_sqlalchemy_engine(
204-
inst_conn_name, user, password, db, resolver=DnsResolver
209+
inst_conn_name, user, password, db, ip_type,resolver=DnsResolver
205210
)
206211

207212
async with pool.connect() as conn:
@@ -217,8 +222,9 @@ async def test_connection_with_asyncpg() -> None:
217222
user = os.environ["POSTGRES_USER"]
218223
password = os.environ["POSTGRES_PASS"]
219224
db = os.environ["POSTGRES_DB"]
225+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
220226

221-
pool, connector = await create_asyncpg_pool(inst_conn_name, user, password, db)
227+
pool, connector = await create_asyncpg_pool(inst_conn_name, user, password, db, ip_type)
222228

223229
async with pool.acquire() as conn:
224230
res = await conn.fetch("SELECT 1")
@@ -233,9 +239,10 @@ async def test_lazy_connection_with_asyncpg() -> None:
233239
user = os.environ["POSTGRES_USER"]
234240
password = os.environ["POSTGRES_PASS"]
235241
db = os.environ["POSTGRES_DB"]
242+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
236243

237244
pool, connector = await create_asyncpg_pool(
238-
inst_conn_name, user, password, db, "lazy"
245+
inst_conn_name, user, password, db, ip_type, "lazy"
239246
)
240247

241248
async with pool.acquire() as conn:

‎tests/system/test_asyncpg_iam_auth.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ async def create_sqlalchemy_engine(
2727
instance_connection_name: str,
2828
user: str,
2929
db: str,
30+
ip_type: str,
3031
refresh_strategy: str = "background",
3132
) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, Connector]:
3233
"""Creates a connection pool for a Cloud SQL instance and returns the pool
@@ -55,6 +56,8 @@ async def create_sqlalchemy_engine(
5556
e.g., [email protected], [email protected]
5657
db (str):
5758
The name of the database, e.g., mydb
59+
ip_type (str):
60+
The IP type of the Cloud SQL instance.
5861
refresh_strategy (Optional[str]):
5962
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6063
or "background". For serverless environments use "lazy" to avoid
@@ -71,9 +74,7 @@ async def create_sqlalchemy_engine(
7174
"asyncpg",
7275
user=user,
7376
db=db,
74-
ip_type=os.environ.get(
75-
"IP_TYPE", "public"
76-
), # can also be "private" or "psc"
77+
ip_type=ip_type,
7778
enable_iam_auth=True,
7879
),
7980
execution_options={"isolation_level": "AUTOCOMMIT"},
@@ -86,8 +87,9 @@ async def test_iam_authn_connection_with_asyncpg() -> None:
8687
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
8788
user = os.environ["POSTGRES_IAM_USER"]
8889
db = os.environ["POSTGRES_DB"]
90+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
8991

90-
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db)
92+
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db, ip_type)
9193

9294
async with pool.connect() as conn:
9395
res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone()
@@ -101,8 +103,9 @@ async def test_lazy_iam_authn_connection_with_asyncpg() -> None:
101103
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
102104
user = os.environ["POSTGRES_IAM_USER"]
103105
db = os.environ["POSTGRES_DB"]
106+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
104107

105-
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db, "lazy")
108+
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db, ip_type, "lazy")
106109

107110
async with pool.connect() as conn:
108111
res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone()

‎tests/system/test_connector_object.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,9 @@ def getconn() -> pymysql.connections.Connection:
3939
user=os.environ["MYSQL_USER"],
4040
password=os.environ["MYSQL_PASS"],
4141
db=os.environ["MYSQL_DB"],
42+
ip_type=os.environ.get(
43+
"IP_TYPE", "public"
44+
),
4245
)
4346
return conn
4447

‎tests/system/test_pg8000_connection.py‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ def create_sqlalchemy_engine(
3232
user: str,
3333
password: str,
3434
db: str,
35+
ip_type: str,
3536
refresh_strategy: str = "background",
3637
resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
3738
) -> tuple[sqlalchemy.engine.Engine, Connector]:
@@ -64,6 +65,8 @@ def create_sqlalchemy_engine(
6465
The database user's password, e.g., secret-password
6566
db (str):
6667
The name of the database, e.g., mydb
68+
ip_type (str):
69+
The IP type of the Cloud SQL instance.
6770
refresh_strategy (Optional[str]):
6871
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6972
or "background". For serverless environments use "lazy" to avoid
@@ -85,9 +88,7 @@ def create_sqlalchemy_engine(
8588
user=user,
8689
password=password,
8790
db=db,
88-
ip_type=os.environ.get(
89-
"IP_TYPE", "public"
90-
), # can also be "private" or "psc"
91+
ip_type=ip_type,
9192
),
9293
)
9394
return engine, connector
@@ -102,8 +103,9 @@ def test_pg8000_connection() -> None:
102103
user = os.environ["POSTGRES_USER"]
103104
password = os.environ["POSTGRES_PASS"]
104105
db = os.environ["POSTGRES_DB"]
106+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
105107

106-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
108+
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db, ip_type)
107109
with engine.connect() as conn:
108110
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
109111
conn.commit()
@@ -118,9 +120,10 @@ def test_lazy_pg8000_connection() -> None:
118120
user = os.environ["POSTGRES_USER"]
119121
password = os.environ["POSTGRES_PASS"]
120122
db = os.environ["POSTGRES_DB"]
123+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
121124

122125
engine, connector = create_sqlalchemy_engine(
123-
inst_conn_name, user, password, db, "lazy"
126+
inst_conn_name, user, password, db, ip_type, "lazy"
124127
)
125128
with engine.connect() as conn:
126129
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
@@ -136,8 +139,9 @@ def test_CAS_pg8000_connection() -> None:
136139
user = os.environ["POSTGRES_USER"]
137140
password = os.environ["POSTGRES_CAS_PASS"]
138141
db = os.environ["POSTGRES_DB"]
142+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
139143

140-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
144+
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db, ip_type)
141145
with engine.connect() as conn:
142146
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
143147
conn.commit()
@@ -152,8 +156,9 @@ def test_customer_managed_CAS_pg8000_connection() -> None:
152156
user = os.environ["POSTGRES_USER"]
153157
password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
154158
db = os.environ["POSTGRES_DB"]
159+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
155160

156-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
161+
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db, ip_type)
157162
with engine.connect() as conn:
158163
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
159164
conn.commit()
@@ -168,9 +173,10 @@ def test_custom_SAN_with_dns_pg8000_connection() -> None:
168173
user = os.environ["POSTGRES_USER"]
169174
password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
170175
db = os.environ["POSTGRES_DB"]
176+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
171177

172178
engine, connector = create_sqlalchemy_engine(
173-
inst_conn_name, user, password, db, resolver=DnsResolver
179+
inst_conn_name, user, password, db, ip_type, resolver=DnsResolver
174180
)
175181
with engine.connect() as conn:
176182
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()

‎tests/system/test_pg8000_iam_auth.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ def create_sqlalchemy_engine(
2626
instance_connection_name: str,
2727
user: str,
2828
db: str,
29+
ip_type: str,
2930
refresh_strategy: str = "background",
3031
) -> tuple[sqlalchemy.engine.Engine, Connector]:
3132
"""Creates a connection pool for a Cloud SQL instance and returns the pool
@@ -55,6 +56,8 @@ def create_sqlalchemy_engine(
5556
e.g., [email protected], [email protected]
5657
db (str):
5758
The name of the database, e.g., mydb
59+
ip_type (str):
60+
The IP type of the Cloud SQL instance.
5861
refresh_strategy (Optional[str]):
5962
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6063
or "background". For serverless environments use "lazy" to avoid
@@ -70,9 +73,7 @@ def create_sqlalchemy_engine(
7073
"pg8000",
7174
user=user,
7275
db=db,
73-
ip_type=os.environ.get(
74-
"IP_TYPE", "public"
75-
), # can also be "private" or "psc"
76+
ip_type=ip_type,
7677
enable_iam_auth=True,
7778
),
7879
)
@@ -84,8 +85,9 @@ def test_pg8000_iam_authn_connection() -> None:
8485
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
8586
user = os.environ["POSTGRES_IAM_USER"]
8687
db = os.environ["POSTGRES_DB"]
88+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
8789

88-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db)
90+
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db, ip_type)
8991
with engine.connect() as conn:
9092
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
9193
conn.commit()
@@ -99,8 +101,9 @@ def test_lazy_pg8000_iam_authn_connection() -> None:
99101
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
100102
user = os.environ["POSTGRES_IAM_USER"]
101103
db = os.environ["POSTGRES_DB"]
104+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
102105

103-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db, "lazy")
106+
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db, ip_type, "lazy")
104107
with engine.connect() as conn:
105108
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
106109
conn.commit()

‎tests/system/test_pymysql_connection.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ def create_sqlalchemy_engine(
2828
user: str,
2929
password: str,
3030
db: str,
31+
ip_type: str,
3132
refresh_strategy: str = "background",
3233
) -> tuple[sqlalchemy.engine.Engine, Connector]:
3334
"""Creates a connection pool for a Cloud SQL instance and returns the pool
@@ -59,6 +60,8 @@ def create_sqlalchemy_engine(
5960
The database user's password, e.g., secret-password
6061
db (str):
6162
The name of the database, e.g., mydb
63+
ip_type (str):
64+
The IP type of the Cloud SQL instance.
6265
refresh_strategy (Optional[str]):
6366
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6467
or "background". For serverless environments use "lazy" to avoid
@@ -75,9 +78,7 @@ def create_sqlalchemy_engine(
7578
user=user,
7679
password=password,
7780
db=db,
78-
ip_type=os.environ.get(
79-
"IP_TYPE", "public"
80-
), # can also be "private" or "psc"
81+
ip_type=ip_type,
8182
),
8283
)
8384
return engine, connector
@@ -92,8 +93,9 @@ def test_pymysql_connection() -> None:
9293
user = os.environ["MYSQL_USER"]
9394
password = os.environ["MYSQL_PASS"]
9495
db = os.environ["MYSQL_DB"]
96+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
9597

96-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
98+
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db, ip_type)
9799
with engine.connect() as conn:
98100
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
99101
conn.commit()
@@ -108,9 +110,10 @@ def test_lazy_pymysql_connection() -> None:
108110
user = os.environ["MYSQL_USER"]
109111
password = os.environ["MYSQL_PASS"]
110112
db = os.environ["MYSQL_DB"]
113+
ip_type = os.environ.get("IP_TYPE", "public") # can be "public", "private" or "psc"
111114

112115
engine, connector = create_sqlalchemy_engine(
113-
inst_conn_name, user, password, db, "lazy"
116+
inst_conn_name, user, password, db, ip_type, "lazy"
114117
)
115118
with engine.connect() as conn:
116119
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()

0 commit comments

Comments
 (0)