from unittest.mock import MagicMock, patch

from feast.infra.compute_engines.spark.utils import _ensure_s3a_event_log_dir

BOTO3_PATH = "feast.infra.compute_engines.spark.utils.boto3"

BOTOCONFIG_PATH = "feast.infra.compute_engines.spark.utils.BotoConfig"

def _base_conf(event_log_dir: str) -> dict:

return {

"spark.eventLog.enabled": "true",

"spark.eventLog.dir": event_log_dir,

"spark.hadoop.fs.s3a.endpoint": "http://minio:9000",

}

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_creates_placeholder_when_empty(mock_boto3):

"""S3A prefix doesn't exist -> placeholder object is written."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 0}

_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/spark-events/"))

s3.list_objects_v2.assert_called_once_with(

Bucket="my-bucket", Prefix="spark-events/", MaxKeys=1

)

s3.put_object.assert_called_once_with(

Bucket="my-bucket", Key="spark-events/.keep", Body=b""

)

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_skips_when_prefix_exists(mock_boto3):

"""S3A prefix already has objects -> no placeholder written."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 3}

_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/spark-events/"))

s3.put_object.assert_not_called()

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_noop_when_event_log_disabled(mock_boto3):

"""spark.eventLog.enabled != true -> boto3 never called."""

_ensure_s3a_event_log_dir(

{"spark.eventLog.enabled": "false", "spark.eventLog.dir": "s3a://b/p/"}

)

mock_boto3.client.assert_not_called()

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_noop_for_non_s3a_path(mock_boto3):

"""Non-S3A paths (hdfs://, file://, etc.) are left untouched."""

_ensure_s3a_event_log_dir(

{"spark.eventLog.enabled": "true", "spark.eventLog.dir": "hdfs:///spark-logs"}

)

mock_boto3.client.assert_not_called()

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_non_fatal_on_s3_error(mock_boto3):

"""boto3 errors are swallowed -> SparkContext will surface its own error."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.side_effect = Exception("connection refused")

_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/spark-events/"))

# ---------------------------------------------------------------------------

# Bucket-root edge cases (s3a://bucket, s3a://bucket/)

# ---------------------------------------------------------------------------

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_bucket_root_no_trailing_slash(mock_boto3):

"""s3a://bucket (no path) -> .keep at bucket root, not /.keep."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 0}

_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket"))

s3.list_objects_v2.assert_called_once_with(Bucket="my-bucket", Prefix="", MaxKeys=1)

s3.put_object.assert_called_once_with(Bucket="my-bucket", Key=".keep", Body=b"")

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_bucket_root_trailing_slash(mock_boto3):

"""s3a://bucket/ (trailing slash, empty prefix) -> .keep at bucket root."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 0}

_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/"))

s3.list_objects_v2.assert_called_once_with(Bucket="my-bucket", Prefix="", MaxKeys=1)

s3.put_object.assert_called_once_with(Bucket="my-bucket", Key=".keep", Body=b"")

# ---------------------------------------------------------------------------

# Credentials from spark config / env var fallback

# ---------------------------------------------------------------------------

@patch.dict(

"os.environ",

{

"AWS_ACCESS_KEY_ID": "env-ak",

"AWS_SECRET_ACCESS_KEY": "env-sk", # pragma: allowlist secret

"AWS_SESSION_TOKEN": "env-st",

},

)

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_uses_spark_config_credentials(mock_boto3):

"""Credentials in spark config take precedence over env vars."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 1}

conf = {

**_base_conf("s3a://my-bucket/logs/"),

"spark.hadoop.fs.s3a.access.key": "spark-ak",

"spark.hadoop.fs.s3a.secret.key": "spark-sk", # pragma: allowlist secret

"spark.hadoop.fs.s3a.session.token": "spark-st",

}

_ensure_s3a_event_log_dir(conf)

mock_boto3.client.assert_called_once()

kw = mock_boto3.client.call_args.kwargs

assert kw["aws_access_key_id"] == "spark-ak"

assert kw["aws_secret_access_key"] == "spark-sk" # pragma: allowlist secret

assert kw["aws_session_token"] == "spark-st"

@patch.dict(

"os.environ",

{

"AWS_ACCESS_KEY_ID": "env-ak",

"AWS_SECRET_ACCESS_KEY": "env-sk", # pragma: allowlist secret

"AWS_SESSION_TOKEN": "env-st",

},

)

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_falls_back_to_env_credentials(mock_boto3):

"""Without spark config keys, env vars are used."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 1}

_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/logs/"))

mock_boto3.client.assert_called_once()

kw = mock_boto3.client.call_args.kwargs

assert kw["aws_access_key_id"] == "env-ak"

assert kw["aws_secret_access_key"] == "env-sk" # pragma: allowlist secret

assert kw["aws_session_token"] == "env-st"

@patch.dict("os.environ", {}, clear=True)

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_no_credentials_passes_none(mock_boto3):

"""No credentials anywhere -> None passed to boto3 (anonymous / instance role)."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 1}

conf = {

"spark.eventLog.enabled": "true",

"spark.eventLog.dir": "s3a://my-bucket/logs/",

}

_ensure_s3a_event_log_dir(conf)

mock_boto3.client.assert_called_once()

kw = mock_boto3.client.call_args.kwargs

assert kw["aws_access_key_id"] is None

assert kw["aws_secret_access_key"] is None

assert kw["aws_session_token"] is None

# ---------------------------------------------------------------------------

# Path-style addressing (MinIO / S3-compatible)

# ---------------------------------------------------------------------------

@patch(BOTOCONFIG_PATH)

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_path_style_when_enabled(mock_boto3, mock_config_cls):

"""spark.hadoop.fs.s3a.path.style.access=true -> addressing_style='path'."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 1}

conf = {

**_base_conf("s3a://my-bucket/logs/"),

"spark.hadoop.fs.s3a.path.style.access": "true",

}

_ensure_s3a_event_log_dir(conf)

mock_config_cls.assert_called_once()

config_kwargs = mock_config_cls.call_args

assert config_kwargs.kwargs["s3"] == {"addressing_style": "path"}

@patch(BOTOCONFIG_PATH)

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_virtual_hosted_style_by_default(

mock_boto3, mock_config_cls

):

"""No path.style.access config -> addressing_style='auto'."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 1}

_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/logs/"))

mock_config_cls.assert_called_once()

config_kwargs = mock_config_cls.call_args

assert config_kwargs.kwargs["s3"] == {"addressing_style": "auto"}

# ---------------------------------------------------------------------------

# Endpoint env var fallback (AWS_ENDPOINT_URL)

# ---------------------------------------------------------------------------

@patch.dict("os.environ", {"AWS_ENDPOINT_URL": "http://localhost:9000"}, clear=True)

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_endpoint_from_env(mock_boto3):

"""AWS_ENDPOINT_URL env var is used when spark config has no endpoint."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 1}

conf = {

"spark.eventLog.enabled": "true",

"spark.eventLog.dir": "s3a://my-bucket/logs/",

}

_ensure_s3a_event_log_dir(conf)

mock_boto3.client.assert_called_once()

kw = mock_boto3.client.call_args.kwargs

assert kw["endpoint_url"] == "http://localhost:9000"

@patch.dict("os.environ", {"AWS_ENDPOINT_URL": "http://env-endpoint:9000"}, clear=True)

@patch(BOTOCONFIG_PATH, MagicMock())

@patch(BOTO3_PATH)

def test_ensure_s3a_event_log_dir_spark_endpoint_over_env(mock_boto3):

"""spark.hadoop.fs.s3a.endpoint takes precedence over AWS_ENDPOINT_URL."""

s3 = MagicMock()

mock_boto3.client.return_value = s3

s3.list_objects_v2.return_value = {"KeyCount": 1}

_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/logs/"))

mock_boto3.client.assert_called_once()

kw = mock_boto3.client.call_args.kwargs

assert kw["endpoint_url"] == "http://minio:9000"