Skip to content

Commit ac96746

Browse files
committed
fix: replace broken S3 cleanup with LlamaStack files API
The S3 file cleanup was failing because boto3 is not a project dependency. Replace the direct S3 access with the LlamaStack client's files API (client.files.list/delete), which requires no additional dependencies. - Remove _cleanup_s3_files helper and its boto3 dependency - Remove S3 credential parameters from distribution fixtures - Add _cleanup_files using the LlamaStack files API - Snapshot existing file IDs before tests to only delete files created during the current test execution
1 parent 8a8185f commit ac96746

File tree

1 file changed

+26
-101
lines changed

1 file changed

+26
-101
lines changed

tests/llama_stack/conftest.py

Lines changed: 26 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -73,59 +73,6 @@
7373
distribution_name = generate_random_name(prefix="llama-stack-distribution")
7474

7575

76-
def _cleanup_s3_files(
77-
bucket_name: str,
78-
endpoint_url: str,
79-
region: str,
80-
access_key_id: str,
81-
secret_access_key: str,
82-
) -> None:
83-
"""
84-
Clean up files from S3 bucket that were uploaded during tests.
85-
86-
Args:
87-
bucket_name: S3 bucket name
88-
endpoint_url: S3 endpoint URL
89-
region: S3 region
90-
access_key_id: AWS access key ID
91-
secret_access_key: AWS secret access key
92-
"""
93-
94-
try:
95-
import boto3
96-
from botocore.exceptions import ClientError
97-
98-
s3_client = boto3.client(
99-
service_name="s3",
100-
endpoint_url=endpoint_url,
101-
aws_access_key_id=access_key_id,
102-
aws_secret_access_key=secret_access_key,
103-
region_name=region,
104-
)
105-
106-
response = s3_client.list_objects_v2(Bucket=bucket_name)
107-
108-
if "Contents" not in response:
109-
LOGGER.info("No files found to clean up from S3")
110-
return
111-
112-
# We only want to delete files that start with "file-"
113-
for obj in response["Contents"]:
114-
key = obj["Key"]
115-
if key.startswith("file-"):
116-
s3_client.delete_object(Bucket=bucket_name, Key=key)
117-
LOGGER.debug(f"Deleted file from S3: {key}")
118-
119-
response = s3_client.list_objects_v2(Bucket=bucket_name)
120-
121-
if "Contents" not in response:
122-
LOGGER.info("No files found to clean up from S3")
123-
return
124-
125-
except ClientError as e:
126-
LOGGER.warning(f"Failed to clean up S3 files: {e}")
127-
128-
12976
@pytest.fixture(scope="class")
13077
def enabled_llama_stack_operator(dsc_resource: DataScienceCluster) -> Generator[DataScienceCluster, Any, Any]:
13178
with update_components_in_dsc(
@@ -385,11 +332,6 @@ def unprivileged_llama_stack_distribution(
385332
enabled_llama_stack_operator: DataScienceCluster,
386333
request: FixtureRequest,
387334
llama_stack_server_config: dict[str, Any],
388-
ci_s3_bucket_name: str,
389-
ci_s3_bucket_endpoint: str,
390-
ci_s3_bucket_region: str,
391-
aws_access_key_id: str,
392-
aws_secret_access_key: str,
393335
unprivileged_llama_stack_distribution_secret: Secret,
394336
unprivileged_postgres_deployment: Deployment,
395337
unprivileged_postgres_service: Service,
@@ -406,25 +348,6 @@ def unprivileged_llama_stack_distribution(
406348
lls_dist.wait_for_status(status=LlamaStackDistribution.Status.READY, timeout=600)
407349
yield lls_dist
408350

409-
try:
410-
env_vars = llama_stack_server_config.get("containerSpec", {}).get("env", [])
411-
enable_s3 = any(env.get("name") == "ENABLE_S3" and env.get("value") == "s3" for env in env_vars)
412-
413-
if enable_s3:
414-
try:
415-
_cleanup_s3_files(
416-
bucket_name=ci_s3_bucket_name,
417-
endpoint_url=ci_s3_bucket_endpoint,
418-
region=ci_s3_bucket_region,
419-
access_key_id=aws_access_key_id,
420-
secret_access_key=aws_secret_access_key,
421-
)
422-
except Exception as e: # noqa: BLE001
423-
LOGGER.warning(f"Failed to clean up S3 files: {e}")
424-
425-
except Exception as e: # noqa: BLE001
426-
LOGGER.warning(f"Failed to clean up S3 files: {e}")
427-
428351

429352
@pytest.fixture(scope="class")
430353
def llama_stack_distribution(
@@ -433,11 +356,6 @@ def llama_stack_distribution(
433356
enabled_llama_stack_operator: DataScienceCluster,
434357
request: FixtureRequest,
435358
llama_stack_server_config: dict[str, Any],
436-
ci_s3_bucket_name: str,
437-
ci_s3_bucket_endpoint: str,
438-
ci_s3_bucket_region: str,
439-
aws_access_key_id: str,
440-
aws_secret_access_key: str,
441359
llama_stack_distribution_secret: Secret,
442360
postgres_deployment: Deployment,
443361
postgres_service: Service,
@@ -453,25 +371,6 @@ def llama_stack_distribution(
453371
lls_dist.wait_for_status(status=LlamaStackDistribution.Status.READY, timeout=600)
454372
yield lls_dist
455373

456-
try:
457-
env_vars = llama_stack_server_config.get("containerSpec", {}).get("env", [])
458-
enable_s3 = any(env.get("name") == "ENABLE_S3" and env.get("value") == "s3" for env in env_vars)
459-
460-
if enable_s3:
461-
try:
462-
_cleanup_s3_files(
463-
bucket_name=ci_s3_bucket_name,
464-
endpoint_url=ci_s3_bucket_endpoint,
465-
region=ci_s3_bucket_region,
466-
access_key_id=aws_access_key_id,
467-
secret_access_key=aws_secret_access_key,
468-
)
469-
except Exception as e: # noqa: BLE001
470-
LOGGER.warning(f"Failed to clean up S3 files: {e}")
471-
472-
except Exception as e: # noqa: BLE001
473-
LOGGER.warning(f"Failed to clean up S3 files: {e}")
474-
475374

476375
def _get_llama_stack_distribution_deployment(
477376
client: DynamicClient,
@@ -642,11 +541,37 @@ def _create_llama_stack_client(
642541
http_client=http_client,
643542
)
644543
wait_for_llama_stack_client_ready(client=client)
544+
existing_file_ids = {f.id for f in client.files.list().data}
545+
645546
yield client
547+
548+
_cleanup_files(client=client, existing_file_ids=existing_file_ids)
646549
finally:
647550
http_client.close()
648551

649552

553+
def _cleanup_files(client: LlamaStackClient, existing_file_ids: set[str]) -> None:
554+
"""Delete files created during test execution via the LlamaStack files API.
555+
556+
Only deletes files whose IDs were not present before the test ran,
557+
avoiding interference with other test sessions.
558+
559+
Args:
560+
client: The LlamaStackClient used during the test
561+
existing_file_ids: File IDs that existed before the test started
562+
"""
563+
try:
564+
for file in client.files.list().data:
565+
if file.id not in existing_file_ids:
566+
try:
567+
client.files.delete(file_id=file.id)
568+
LOGGER.debug(f"Deleted file: {file.id}")
569+
except Exception as e:
570+
LOGGER.warning(f"Failed to delete file {file.id}: {e}")
571+
except Exception as e:
572+
LOGGER.warning(f"Failed to clean up files: {e}")
573+
574+
650575
@pytest.fixture(scope="class")
651576
def unprivileged_llama_stack_client(
652577
unprivileged_llama_stack_test_route: Route,

0 commit comments

Comments
 (0)