Skip to content

Commit 63c5eb1

Browse files
ntkatholentkathole
committed
fix: Fix intermittent async test failures for DynamoDB and Redis
DynamoDB (integration tests): - Add conftest.py with dynamodb_local_environment fixture that creates a fully isolated Environment (its own DynamoDB Local container + FileDataSourceCreator) with dummy AWS credentials set *before* any boto client is instantiated. This prevents the expired STS session token present in CI from bleeding into the lazy aiobotocore client and causing UnrecognizedClientException on DynamoDB Local 2.x. - test_push_features_and_read_async: scope to mongodb only via the shared environment; add test_push_features_and_read_async_dynamodb that uses the isolated fixture for the same coverage. - test_async_online_retrieval_with_event_timestamps_dynamo: replace shared environment + broken monkeypatch fixture with the isolated dynamodb_local_environment fixture. Redis (unit test): - Replace asyncio.get_event_loop().run_until_complete() with asyncio.run() in test_online_write_batch_async_skip_dedup_single_pipeline to fix RuntimeError: There is no current event loop in thread 'MainThread' on Python 3.10+. Signed-off-by: ntkathole <nikhilkathole2683@gmail.com>
1 parent a4fde83 commit 63c5eb1

5 files changed

Lines changed: 220 additions & 5 deletions

File tree

‎infra/feast-operator/bundle/manifests/openlineage-secret_v1_secret.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,4 @@ kind: Secret
33
metadata:
44
name: openlineage-secret
55
stringData:
6-
api_key: your-marquez-api-key # pragma: allowlist secret
6+
api_key: your-marquez-api-key #pragma: allowlist secret

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

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
"""
2+
Fixtures for online-store integration tests.
3+
"""
4+
5+
from typing import Dict
6+
7+
import pytest
8+
from testcontainers.core.container import DockerContainer
9+
from testcontainers.core.waiting_utils import wait_for_logs
10+
11+
from tests.universal.feature_repos.universal.online_store_creator import (
12+
OnlineStoreCreator,
13+
)
14+
15+
16+
class _SharedDbDynamoDBOnlineStoreCreator(OnlineStoreCreator):
17+
"""DynamoDB Local container started with ``-sharedDb -inMemory``.
18+
19+
Why ``-sharedDb``
20+
-----------------
21+
DynamoDB Local 2.x namespaces tables by the **access key ID** in the
22+
request signature. In CI, the sync ``boto3`` client and the async
23+
``aiobotocore`` client can resolve credentials from *different* sources
24+
(env vars, credential file, ``credential_process``, container IAM role,
25+
etc.) even after ``monkeypatch.setenv`` has set fake keys—because the
26+
credential chain is evaluated lazily and various caches may hold stale
27+
values.
28+
29+
When the two clients end up using *different* access keys, the sync
30+
client creates tables in namespace A while the async client queries
31+
namespace B, which is empty → ``ResourceNotFoundException``.
32+
33+
``-sharedDb`` collapses all namespaces into a single in-memory database,
34+
making table visibility completely independent of which credentials each
35+
client uses. This is the correct setting for integration tests that want
36+
to verify async read/write behaviour without caring about credential
37+
isolation.
38+
"""
39+
40+
def __init__(self, project_name: str, **kwargs):
41+
super().__init__(project_name)
42+
self.container = (
43+
DockerContainer("amazon/dynamodb-local:latest")
44+
.with_exposed_ports("8000")
45+
.with_command("-jar DynamoDBLocal.jar -sharedDb -inMemory")
46+
)
47+
48+
def create_online_store(self) -> Dict[str, str]:
49+
self.container.start()
50+
wait_for_logs(
51+
container=self.container,
52+
predicate="Initializing DynamoDB Local with the following configuration:",
53+
timeout=10,
54+
)
55+
exposed_port = self.container.get_exposed_port("8000")
56+
return {
57+
"type": "dynamodb",
58+
"endpoint_url": f"http://localhost:{exposed_port}",
59+
"region": "us-west-2",
60+
}
61+
62+
def teardown(self):
63+
self.container.stop()
64+
65+
66+
@pytest.fixture
67+
async def dynamodb_local_environment(monkeypatch, worker_id):
68+
"""Isolated, self-contained Environment for DynamoDB async tests.
69+
70+
Root cause of the async credential failures
71+
-------------------------------------------
72+
DynamoDB Local 2.x isolates tables **per access key ID**. In CI,
73+
``boto3`` (sync, used to provision tables via ``store.apply()``) and
74+
``aiobotocore`` (async, used for reads/writes in the test body) may
75+
resolve credentials from *different* sources even when ``monkeypatch``
76+
has set fake static keys—the credential chain is evaluated lazily and
77+
caches may hold stale values from a real AWS session configured in the
78+
runner environment.
79+
80+
When the two clients end up using different access key IDs they land in
81+
different DynamoDB Local namespaces:
82+
83+
* sync client → namespace ``KEY_A`` → tables exist ✓
84+
* async client → namespace ``KEY_B`` → tables not found → ``ResourceNotFoundException``
85+
86+
Fix: ``_SharedDbDynamoDBOnlineStoreCreator``
87+
--------------------------------------------
88+
The isolated container is started with ``-sharedDb -inMemory``. In
89+
shared-DB mode DynamoDB Local stores *all* tables in a single namespace
90+
regardless of the access key, so sync and async clients always see the
91+
same tables.
92+
93+
Why async + ``await fs.initialize()`` before yielding
94+
-----------------------------------------------------
95+
Calling ``await fs.initialize()`` eagerly creates the ``aiobotocore``
96+
client inside this fixture's event loop (the *same* loop the test will
97+
run in). This pre-caches:
98+
99+
1. ``FeatureStore._provider`` so the identical ``DynamoDBOnlineStore``
100+
instance is reused for the entire test.
101+
2. The aiobotocore client, which is now unambiguously pointed at our
102+
isolated container's ``endpoint_url``.
103+
104+
Yields
105+
------
106+
tuple[Environment, TestData]
107+
``(environment, (entities, datasets, data_sources))``
108+
"""
109+
from feast.infra.online_stores.dynamodb import DynamoDBOnlineStore
110+
from tests.universal.feature_repos.integration_test_repo_config import (
111+
IntegrationTestRepoConfig,
112+
)
113+
from tests.universal.feature_repos.repo_configuration import (
114+
construct_test_environment,
115+
construct_universal_test_data,
116+
)
117+
from tests.universal.feature_repos.universal.data_sources.file import (
118+
FileDataSourceCreator,
119+
)
120+
121+
# Set fake static credentials before any boto client is created.
122+
# These are accepted by DynamoDB Local regardless of validity.
123+
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "fakeaccesskey000000")
124+
monkeypatch.setenv(
125+
"AWS_SECRET_ACCESS_KEY", "fakesecretkey0000000000000000000000000000"
126+
)
127+
monkeypatch.delenv("AWS_SESSION_TOKEN", raising=False)
128+
monkeypatch.delenv("AWS_SECURITY_TOKEN", raising=False)
129+
# Prevent IMDS from injecting real session tokens on EC2-backed runners.
130+
monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true")
131+
# Disable the container credentials provider (ECS/EKS IAM roles).
132+
monkeypatch.delenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", raising=False)
133+
monkeypatch.delenv("AWS_CONTAINER_CREDENTIALS_FULL_URI", raising=False)
134+
# Ensure no profile redirects boto to a different credential source.
135+
monkeypatch.delenv("AWS_PROFILE", raising=False)
136+
monkeypatch.delenv("AWS_DEFAULT_PROFILE", raising=False)
137+
138+
# Reset class-level boto3 client caches so that no stale client from a
139+
# previous test in this worker bleeds into our isolated environment.
140+
DynamoDBOnlineStore._dynamodb_client = None
141+
DynamoDBOnlineStore._dynamodb_resource = None
142+
143+
config = IntegrationTestRepoConfig(
144+
provider="local",
145+
offline_store_creator=FileDataSourceCreator,
146+
online_store_creator=_SharedDbDynamoDBOnlineStoreCreator,
147+
online_store=None,
148+
)
149+
150+
environment = construct_test_environment(
151+
config,
152+
fixture_request=None,
153+
worker_id=worker_id,
154+
)
155+
environment.setup()
156+
157+
# FileDataSourceCreator writes only local Parquet files — no AWS calls.
158+
universal_test_data = construct_universal_test_data(environment)
159+
160+
# Eagerly initialise the aiobotocore client in *this* event loop so it
161+
# is guaranteed to point at our container and is reused throughout the
162+
# test body without lazy-init surprises.
163+
await environment.feature_store.initialize()
164+
165+
yield environment, universal_test_data
166+
167+
# Cleanly shut down the async client before the container disappears.
168+
await environment.feature_store.close()
169+
environment.teardown()
170+
171+
# Flush class-level caches so the next test starts completely fresh.
172+
DynamoDBOnlineStore._dynamodb_client = None
173+
DynamoDBOnlineStore._dynamodb_resource = None

‎sdk/python/tests/integration/online_store/test_push_features_to_online_store.py‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def test_push_features_and_read(store):
4747

4848

4949
@pytest.mark.integration
50-
@pytest.mark.universal_online_stores(only=["dynamodb", "mongodb"])
50+
@pytest.mark.universal_online_stores(only=["mongodb"])
5151
async def test_push_features_and_read_async(store):
5252
await store.push_async("location_stats_push_source", _ingest_df())
5353

@@ -56,3 +56,36 @@ async def test_push_features_and_read_async(store):
5656
entity_rows=[{"location_id": 1}],
5757
)
5858
assert_response(online_resp)
59+
60+
61+
@pytest.mark.asyncio
62+
@pytest.mark.integration
63+
@pytest.mark.universal_online_stores
64+
async def test_push_features_and_read_async_dynamodb(dynamodb_local_environment):
65+
"""Async push + async read for DynamoDB with a credential-isolated environment.
66+
67+
DynamoDB Local 2.x rejects requests that carry an expired AWS session
68+
token. In CI, real (possibly expired) STS credentials exist in the
69+
environment. The shared ``environment`` fixture resolves credentials
70+
before the async client is created, so those bad credentials bleed in.
71+
72+
This test uses ``dynamodb_local_environment``, which sets dummy
73+
credentials *before* any boto client is instantiated, guaranteeing that
74+
both the sync boto3 table-provisioning client and the async aiobotocore
75+
client start with clean, token-free credentials.
76+
"""
77+
environment, universal_test_data = dynamodb_local_environment
78+
store = environment.feature_store
79+
_, _, data_sources = universal_test_data
80+
81+
feature_views = construct_universal_feature_views(data_sources)
82+
location_fv = feature_views.pushed_locations
83+
store.apply([location(), location_fv])
84+
85+
await store.push_async("location_stats_push_source", _ingest_df())
86+
87+
online_resp = await store.get_online_features_async(
88+
features=["pushable_location_stats:temperature"],
89+
entity_rows=[{"location_id": 1}],
90+
)
91+
assert_response(online_resp)

‎sdk/python/tests/integration/online_store/test_universal_online.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -532,10 +532,19 @@ async def test_async_online_retrieval_with_event_timestamps(
532532

533533
@pytest.mark.asyncio
534534
@pytest.mark.integration
535-
@pytest.mark.universal_online_stores(only=["dynamodb"])
535+
@pytest.mark.universal_online_stores
536536
async def test_async_online_retrieval_with_event_timestamps_dynamo(
537-
environment, universal_data_sources
537+
dynamodb_local_environment,
538538
):
539+
"""Async online retrieval for DynamoDB with a credential-isolated environment.
540+
541+
Uses ``dynamodb_local_environment`` (its own DynamoDB Local container +
542+
FileDataSourceCreator) so that dummy credentials are set before any boto
543+
client is created. This avoids the expired-STS-token problem that
544+
occurs when aiobotocore lazily resolves credentials from the shared
545+
environment in CI.
546+
"""
547+
environment, universal_data_sources = dynamodb_local_environment
539548
await _do_async_retrieval_test(environment, universal_data_sources)
540549

541550

‎sdk/python/tests/unit/infra/online_store/test_redis.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -468,7 +468,7 @@ async def _run():
468468
config, feature_view, data, progress=None
469469
)
470470

471-
asyncio.get_event_loop().run_until_complete(_run())
471+
asyncio.run(_run())
472472

473473
assert mock_async_client.pipeline.call_count == 1
474474
async_pipe.hmget.assert_not_called()

0 commit comments

Comments
 (0)