From 207b7c70c0ca8de3926e04e1e19fb5ee10e5580c Mon Sep 17 00:00:00 2001
From: Justus Schock <12886177+justusschock@users.noreply.github.com>
Date: Thu, 13 Jul 2023 19:07:29 +0200
Subject: [PATCH] Fix IAM Credentials Backend (#18073)

* update

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update

* update

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update

* test-updates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test-updates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 src/lightning/data/backends.py    |  4 +++-
 tests/tests_data/test_backends.py | 33 +++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 tests/tests_data/test_backends.py

diff --git a/src/lightning/data/backends.py b/src/lightning/data/backends.py
index 9cdc764185..1d4dbcccd4 100644
--- a/src/lightning/data/backends.py
+++ b/src/lightning/data/backends.py
@@ -47,7 +47,9 @@ class S3DatasetBackend:
         if os.getenv("AWS_ACCESS_KEY") and os.getenv("AWS_SECRET_ACCESS_KEY"):
             return {"access_key": os.getenv("AWS_ACCESS_KEY"), "secret_key": os.getenv("AWS_SECRET_ACCESS_KEY")}
 
-        return self.get_aws_credentials()
+        aws_creds = self.get_aws_credentials()
+
+        return {"access_key": aws_creds.access_key, "secret_key": aws_creds.secret_key, "token": aws_creds.token}
 
     def handle_error(self, exc: Exception) -> None:
         from botocore.exceptions import NoCredentialsError
diff --git a/tests/tests_data/test_backends.py b/tests/tests_data/test_backends.py
new file mode 100644
index 0000000000..d15f666f43
--- /dev/null
+++ b/tests/tests_data/test_backends.py
@@ -0,0 +1,33 @@
+import os
+from collections import namedtuple
+from typing import Mapping
+from unittest import mock
+
+
+def test_s3_dataset_backend_credentials_env_vars():
+    from lightning.data.backends import S3DatasetBackend
+
+    os.environ["AWS_ACCESS_KEY"] = "123"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "abc"
+
+    assert S3DatasetBackend().credentials() == {"access_key": "123", "secret_key": "abc"}
+    os.environ.pop("AWS_ACCESS_KEY")
+    os.environ.pop("AWS_SECRET_ACCESS_KEY")
+
+
+_Credentials = namedtuple("RefreshableCredentials", ("access_key", "secret_key", "token"))
+
+
+@mock.patch("botocore.credentials.InstanceMetadataProvider.load", return_value=_Credentials("abc", "def", "ghi"))
+def test_s3_dataset_backend_credentials_iam(patch1):
+    from lightning.data.backends import S3DatasetBackend
+
+    credentials = S3DatasetBackend().credentials()
+    assert isinstance(credentials, Mapping)
+    assert credentials == {"access_key": "abc", "secret_key": "def", "token": "ghi"}
+
+
+def test_local_dataset_backend_credentials():
+    from lightning.data.backends import LocalDatasetBackend
+
+    assert LocalDatasetBackend().credentials() == {}