Skip to content

Commit

Permalink
fix(storage-test): Use access token in real s3 tests
Browse files Browse the repository at this point in the history
  • Loading branch information
rahulinux committed Dec 11, 2024
1 parent a6c8832 commit b8db228
Showing 1 changed file with 13 additions and 3 deletions.
16 changes: 13 additions & 3 deletions test_runner/fixtures/remote_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,9 @@ def access_key(self) -> str:
def secret_key(self) -> str:
return "test"

def session_token(self) -> str:
return "test"

def kill(self):
self.server.stop()

Expand Down Expand Up @@ -161,6 +164,7 @@ class S3Storage:
bucket_region: str
access_key: str | None
secret_key: str | None
session_token: str | None
aws_profile: str | None
prefix_in_bucket: str
client: S3Client
Expand All @@ -181,13 +185,14 @@ def access_env_vars(self) -> dict[str, str]:
if home is not None:
env["HOME"] = home
return env
if self.access_key is not None and self.secret_key is not None:
if self.access_key is not None and self.secret_key is not None and self.session_token is not None:
return {
"AWS_ACCESS_KEY_ID": self.access_key,
"AWS_SECRET_ACCESS_KEY": self.secret_key,
"AWS_SESSION_TOKEN": self.session_token,
}
raise RuntimeError(
"Either AWS_PROFILE or (AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY) have to be set for S3Storage"
"Either AWS_PROFILE or (AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY and AWS_SESSION_TOKEN) have to be set for S3Storage"
)

def to_string(self) -> str:
Expand Down Expand Up @@ -352,13 +357,15 @@ def to_bucket_name(user: str, test_name: str) -> str:
mock_region = mock_s3_server.region()

access_key, secret_key = mock_s3_server.access_key(), mock_s3_server.secret_key()
session_token = mock_s3_server.session_token()

client = boto3.client(
"s3",
endpoint_url=mock_endpoint,
region_name=mock_region,
aws_access_key_id=access_key,
aws_secret_access_key=secret_key,
aws_access_token=session_token,
)

bucket_name = to_bucket_name(user, test_name)
Expand All @@ -372,6 +379,7 @@ def to_bucket_name(user: str, test_name: str) -> str:
bucket_region=mock_region,
access_key=access_key,
secret_key=secret_key,
session_token=session_token,
aws_profile=None,
prefix_in_bucket="",
client=client,
Expand All @@ -383,9 +391,10 @@ def to_bucket_name(user: str, test_name: str) -> str:

env_access_key = os.getenv("AWS_ACCESS_KEY_ID")
env_secret_key = os.getenv("AWS_SECRET_ACCESS_KEY")
env_access_token = os.getenv("AWS_SESSION_TOKEN")
env_profile = os.getenv("AWS_PROFILE")
assert (
env_access_key and env_secret_key
env_access_key and env_secret_key and env_access_token
) or env_profile, "need to specify either access key and secret access key or profile"

bucket_name = bucket_name or os.getenv("REMOTE_STORAGE_S3_BUCKET")
Expand All @@ -405,6 +414,7 @@ def to_bucket_name(user: str, test_name: str) -> str:
bucket_region=bucket_region,
access_key=env_access_key,
secret_key=env_secret_key,
session_token=env_access_token,
aws_profile=env_profile,
prefix_in_bucket=prefix_in_bucket,
client=client,
Expand Down

0 comments on commit b8db228

Please sign in to comment.