Skip to content

Commit da24656

Browse files
tokokotokoko
authored
chore: Scope fixtures to session in test_universal_registry (#4497)
chore: refactor fixtures in test_universal_registry Signed-off-by: tokoko <togurgenidze@gmail.com>
1 parent 96344b2 commit da24656

2 files changed

Lines changed: 147 additions & 120 deletions

File tree

‎sdk/python/tests/integration/conftest.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22

33
import pytest
44
from testcontainers.keycloak import KeycloakContainer
5+
from testcontainers.minio import MinioContainer
6+
from testcontainers.mysql import MySqlContainer
7+
from testcontainers.postgres import PostgresContainer
58

69
from tests.utils.auth_permissions_util import setup_permissions_on_keycloak
710

@@ -14,3 +17,33 @@ def start_keycloak_server():
1417
with KeycloakContainer("quay.io/keycloak/keycloak:24.0.1") as keycloak_container:
1518
setup_permissions_on_keycloak(keycloak_container.get_client())
1619
yield keycloak_container.get_url()
20+
21+
22+
@pytest.fixture(scope="session")
23+
def mysql_server():
24+
container = MySqlContainer("mysql:latest")
25+
container.start()
26+
27+
yield container
28+
29+
container.stop()
30+
31+
32+
@pytest.fixture(scope="session")
33+
def postgres_server():
34+
container = PostgresContainer()
35+
container.start()
36+
37+
yield container
38+
39+
container.stop()
40+
41+
42+
@pytest.fixture(scope="session")
43+
def minio_server():
44+
container = MinioContainer()
45+
container.start()
46+
47+
yield container
48+
49+
container.stop()

‎sdk/python/tests/integration/registration/test_universal_registry.py‎

Lines changed: 114 additions & 120 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
# limitations under the License.
1414
import logging
1515
import os
16+
import random
17+
import string
1618
import time
1719
from datetime import timedelta, timezone
1820
from tempfile import mkstemp
@@ -22,10 +24,8 @@
2224
import pandas as pd
2325
import pytest
2426
from pytest_lazyfixture import lazy_fixture
25-
from testcontainers.core.container import DockerContainer
26-
from testcontainers.core.waiting_utils import wait_for_logs
27-
from testcontainers.minio import MinioContainer
2827
from testcontainers.mysql import MySqlContainer
28+
from testcontainers.postgres import PostgresContainer
2929

3030
from feast import FeatureService, FileSource, RequestSource
3131
from feast.data_format import AvroFormat, ParquetFormat
@@ -93,158 +93,136 @@ def s3_registry() -> Registry:
9393

9494

9595
@pytest.fixture(scope="function")
96-
def minio_registry() -> Registry:
97-
bucket_name = "test-bucket"
96+
def minio_registry(minio_server):
97+
bucket_name = "".join(random.choices(string.ascii_lowercase, k=10))
9898

99-
container = MinioContainer()
100-
container.start()
101-
client = container.get_client()
99+
client = minio_server.get_client()
102100
client.make_bucket(bucket_name)
103101

104-
container_host = container.get_container_host_ip()
105-
exposed_port = container.get_exposed_port(container.port)
102+
container_host = minio_server.get_container_host_ip()
103+
exposed_port = minio_server.get_exposed_port(minio_server.port)
106104

107105
registry_config = RegistryConfig(
108106
path=f"s3://{bucket_name}/registry.db", cache_ttl_seconds=600
109107
)
110108

111109
mock_environ = {
112110
"FEAST_S3_ENDPOINT_URL": f"http://{container_host}:{exposed_port}",
113-
"AWS_ACCESS_KEY_ID": container.access_key,
114-
"AWS_SECRET_ACCESS_KEY": container.secret_key,
111+
"AWS_ACCESS_KEY_ID": minio_server.access_key,
112+
"AWS_SECRET_ACCESS_KEY": minio_server.secret_key,
115113
"AWS_SESSION_TOKEN": "",
116114
}
117115

118116
with mock.patch.dict(os.environ, mock_environ):
119117
yield Registry("project", registry_config, None)
120118

121-
container.stop()
122-
123-
124-
POSTGRES_USER = "test"
125-
POSTGRES_PASSWORD = "test"
126-
POSTGRES_DB = "test"
127119

128120
logger = logging.getLogger(__name__)
129121

130122

131123
@pytest.fixture(scope="function")
132-
def pg_registry():
133-
container = (
134-
DockerContainer("postgres:latest")
135-
.with_exposed_ports(5432)
136-
.with_env("POSTGRES_USER", POSTGRES_USER)
137-
.with_env("POSTGRES_PASSWORD", POSTGRES_PASSWORD)
138-
.with_env("POSTGRES_DB", POSTGRES_DB)
139-
)
140-
141-
container.start()
142-
143-
registry_config = _given_registry_config_for_pg_sql(container)
144-
145-
yield SqlRegistry(registry_config, "project", None)
124+
def pg_registry(postgres_server):
125+
db_name = "".join(random.choices(string.ascii_lowercase, k=10))
146126

147-
container.stop()
127+
_create_pg_database(postgres_server, db_name)
148128

129+
container_port = postgres_server.get_exposed_port(5432)
130+
container_host = postgres_server.get_container_host_ip()
149131

150-
@pytest.fixture(scope="function")
151-
def pg_registry_async():
152-
container = (
153-
DockerContainer("postgres:latest")
154-
.with_exposed_ports(5432)
155-
.with_env("POSTGRES_USER", POSTGRES_USER)
156-
.with_env("POSTGRES_PASSWORD", POSTGRES_PASSWORD)
157-
.with_env("POSTGRES_DB", POSTGRES_DB)
132+
registry_config = SqlRegistryConfig(
133+
registry_type="sql",
134+
cache_ttl_seconds=2,
135+
cache_mode="sync",
136+
# The `path` must include `+psycopg` in order for `sqlalchemy.create_engine()`
137+
# to understand that we are using psycopg3.
138+
path=f"postgresql+psycopg://{postgres_server.username}:{postgres_server.password}@{container_host}:{container_port}/{db_name}",
139+
sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
140+
thread_pool_executor_worker_count=0,
141+
purge_feast_metadata=False,
158142
)
159143

160-
container.start()
161-
162-
registry_config = _given_registry_config_for_pg_sql(container, 2, "thread", 3)
163-
164144
yield SqlRegistry(registry_config, "project", None)
165145

166-
container.stop()
167146

147+
@pytest.fixture(scope="function")
148+
def pg_registry_async(postgres_server):
149+
db_name = "".join(random.choices(string.ascii_lowercase, k=10))
150+
151+
_create_pg_database(postgres_server, db_name)
168152

169-
def _given_registry_config_for_pg_sql(
170-
container,
171-
cache_ttl_seconds=2,
172-
cache_mode="sync",
173-
thread_pool_executor_worker_count=0,
174-
purge_feast_metadata=False,
175-
):
176-
log_string_to_wait_for = "database system is ready to accept connections"
177-
waited = wait_for_logs(
178-
container=container,
179-
predicate=log_string_to_wait_for,
180-
timeout=30,
181-
interval=10,
182-
)
183-
logger.info("Waited for %s seconds until postgres container was up", waited)
184-
container_port = container.get_exposed_port(5432)
185-
container_host = container.get_container_host_ip()
153+
container_port = postgres_server.get_exposed_port(5432)
154+
container_host = postgres_server.get_container_host_ip()
186155

187-
return SqlRegistryConfig(
156+
registry_config = SqlRegistryConfig(
188157
registry_type="sql",
189-
cache_ttl_seconds=cache_ttl_seconds,
190-
cache_mode=cache_mode,
158+
cache_ttl_seconds=2,
159+
cache_mode="thread",
191160
# The `path` must include `+psycopg` in order for `sqlalchemy.create_engine()`
192161
# to understand that we are using psycopg3.
193-
path=f"postgresql+psycopg://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{container_host}:{container_port}/{POSTGRES_DB}",
162+
path=f"postgresql+psycopg://{postgres_server.username}:{postgres_server.password}@{container_host}:{container_port}/{db_name}",
194163
sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
195-
thread_pool_executor_worker_count=thread_pool_executor_worker_count,
196-
purge_feast_metadata=purge_feast_metadata,
164+
thread_pool_executor_worker_count=3,
165+
purge_feast_metadata=False,
197166
)
198167

168+
yield SqlRegistry(registry_config, "project", None)
199169

200-
@pytest.fixture(scope="function")
201-
def mysql_registry():
202-
container = MySqlContainer("mysql:latest")
203-
container.start()
204170

205-
registry_config = _given_registry_config_for_mysql(container)
171+
def _create_mysql_database(container: MySqlContainer, database: str):
172+
container.exec(
173+
f"mysql -uroot -p{container.root_password} -e 'CREATE DATABASE {database}; GRANT ALL PRIVILEGES ON {database}.* TO {container.username};'"
174+
)
206175

207-
yield SqlRegistry(registry_config, "project", None)
208176

209-
container.stop()
177+
def _create_pg_database(container: PostgresContainer, database: str):
178+
container.exec(f"psql -U {container.username} -c 'CREATE DATABASE {database}'")
210179

211180

212181
@pytest.fixture(scope="function")
213-
def mysql_registry_async():
214-
container = MySqlContainer("mysql:latest")
215-
container.start()
182+
def mysql_registry(mysql_server):
183+
db_name = "".join(random.choices(string.ascii_lowercase, k=10))
184+
185+
_create_mysql_database(mysql_server, db_name)
216186

217-
registry_config = _given_registry_config_for_mysql(container, 2, "thread", 3)
187+
connection_url = (
188+
"/".join(mysql_server.get_connection_url().split("/")[:-1]) + f"/{db_name}"
189+
)
190+
191+
registry_config = SqlRegistryConfig(
192+
registry_type="sql",
193+
path=connection_url,
194+
cache_ttl_seconds=2,
195+
cache_mode="sync",
196+
sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
197+
thread_pool_executor_worker_count=0,
198+
purge_feast_metadata=False,
199+
)
218200

219201
yield SqlRegistry(registry_config, "project", None)
220202

221-
container.stop()
222203

204+
@pytest.fixture(scope="function")
205+
def mysql_registry_async(mysql_server):
206+
db_name = "".join(random.choices(string.ascii_lowercase, k=10))
223207

224-
def _given_registry_config_for_mysql(
225-
container,
226-
cache_ttl_seconds=2,
227-
cache_mode="sync",
228-
thread_pool_executor_worker_count=0,
229-
purge_feast_metadata=False,
230-
):
231-
import sqlalchemy
208+
_create_mysql_database(mysql_server, db_name)
232209

233-
engine = sqlalchemy.create_engine(
234-
container.get_connection_url(), pool_pre_ping=True
210+
connection_url = (
211+
"/".join(mysql_server.get_connection_url().split("/")[:-1]) + f"/{db_name}"
235212
)
236-
engine.connect()
237213

238-
return SqlRegistryConfig(
214+
registry_config = SqlRegistryConfig(
239215
registry_type="sql",
240-
path=container.get_connection_url(),
241-
cache_ttl_seconds=cache_ttl_seconds,
242-
cache_mode=cache_mode,
216+
path=connection_url,
217+
cache_ttl_seconds=2,
218+
cache_mode="thread",
243219
sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
244-
thread_pool_executor_worker_count=thread_pool_executor_worker_count,
245-
purge_feast_metadata=purge_feast_metadata,
220+
thread_pool_executor_worker_count=3,
221+
purge_feast_metadata=False,
246222
)
247223

224+
yield SqlRegistry(registry_config, "project", None)
225+
248226

249227
@pytest.fixture(scope="session")
250228
def sqlite_registry():
@@ -339,11 +317,11 @@ def mock_remote_registry():
339317
async_sql_fixtures = [
340318
pytest.param(
341319
lazy_fixture("pg_registry_async"),
342-
marks=pytest.mark.xdist_group(name="pg_registry_async"),
320+
marks=pytest.mark.xdist_group(name="pg_registry"),
343321
),
344322
pytest.param(
345323
lazy_fixture("mysql_registry_async"),
346-
marks=pytest.mark.xdist_group(name="mysql_registry_async"),
324+
marks=pytest.mark.xdist_group(name="mysql_registry"),
347325
),
348326
]
349327

@@ -1609,45 +1587,61 @@ def local_registry_purge_feast_metadata() -> Registry:
16091587

16101588

16111589
@pytest.fixture(scope="function")
1612-
def pg_registry_purge_feast_metadata():
1613-
container = (
1614-
DockerContainer("postgres:latest")
1615-
.with_exposed_ports(5432)
1616-
.with_env("POSTGRES_USER", POSTGRES_USER)
1617-
.with_env("POSTGRES_PASSWORD", POSTGRES_PASSWORD)
1618-
.with_env("POSTGRES_DB", POSTGRES_DB)
1619-
)
1590+
def pg_registry_purge_feast_metadata(postgres_server):
1591+
db_name = "".join(random.choices(string.ascii_lowercase, k=10))
16201592

1621-
container.start()
1593+
_create_pg_database(postgres_server, db_name)
16221594

1623-
registry_config = _given_registry_config_for_pg_sql(container, 2, "thread", 3, True)
1595+
container_port = postgres_server.get_exposed_port(5432)
1596+
container_host = postgres_server.get_container_host_ip()
16241597

1625-
yield SqlRegistry(registry_config, "project", None)
1598+
registry_config = SqlRegistryConfig(
1599+
registry_type="sql",
1600+
cache_ttl_seconds=2,
1601+
cache_mode="thread",
1602+
# The `path` must include `+psycopg` in order for `sqlalchemy.create_engine()`
1603+
# to understand that we are using psycopg3.
1604+
path=f"postgresql+psycopg://{postgres_server.username}:{postgres_server.password}@{container_host}:{container_port}/{db_name}",
1605+
sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
1606+
thread_pool_executor_worker_count=3,
1607+
purge_feast_metadata=True,
1608+
)
16261609

1627-
container.stop()
1610+
yield SqlRegistry(registry_config, "project", None)
16281611

16291612

16301613
@pytest.fixture(scope="function")
1631-
def mysql_registry_purge_feast_metadata():
1632-
container = MySqlContainer("mysql:latest")
1633-
container.start()
1614+
def mysql_registry_purge_feast_metadata(mysql_server):
1615+
db_name = "".join(random.choices(string.ascii_lowercase, k=10))
16341616

1635-
registry_config = _given_registry_config_for_mysql(container, 2, "thread", 3, True)
1617+
_create_mysql_database(mysql_server, db_name)
16361618

1637-
yield SqlRegistry(registry_config, "project", None)
1619+
connection_url = (
1620+
"/".join(mysql_server.get_connection_url().split("/")[:-1]) + f"/{db_name}"
1621+
)
16381622

1639-
container.stop()
1623+
registry_config = SqlRegistryConfig(
1624+
registry_type="sql",
1625+
path=connection_url,
1626+
cache_ttl_seconds=2,
1627+
cache_mode="thread",
1628+
sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
1629+
thread_pool_executor_worker_count=3,
1630+
purge_feast_metadata=True,
1631+
)
1632+
1633+
yield SqlRegistry(registry_config, "project", None)
16401634

16411635

16421636
purge_feast_metadata_fixtures = [
16431637
lazy_fixture("local_registry_purge_feast_metadata"),
16441638
pytest.param(
16451639
lazy_fixture("pg_registry_purge_feast_metadata"),
1646-
marks=pytest.mark.xdist_group(name="pg_registry_purge_feast_metadata"),
1640+
marks=pytest.mark.xdist_group(name="pg_registry"),
16471641
),
16481642
pytest.param(
16491643
lazy_fixture("mysql_registry_purge_feast_metadata"),
1650-
marks=pytest.mark.xdist_group(name="mysql_registry_purge_feast_metadata"),
1644+
marks=pytest.mark.xdist_group(name="mysql_registry"),
16511645
),
16521646
]
16531647

0 commit comments

Comments
 (0)