From 6f064b10666c45b12282d7cc26aa9d7dd78a7b9a Mon Sep 17 00:00:00 2001
From: Francisco Rivera <franmrivera@gmail.com>
Date: Wed, 4 Oct 2023 09:53:28 -0300
Subject: [PATCH] refresh-credentials command

---
 leverage/container.py                   |  5 +++++
 leverage/modules/terraform.py           |  9 +++++++++
 tests/test_containers/__init__.py       |  1 +
 tests/test_containers/test_terraform.py | 11 +++++++++++
 4 files changed, 26 insertions(+)

diff --git a/leverage/container.py b/leverage/container.py
index af9f7c2a..35c948b3 100644
--- a/leverage/container.py
+++ b/leverage/container.py
@@ -661,6 +661,11 @@ def check_for_layer_location(self):
             logger.error("This command can only run at [bold]layer[/bold] level.")
             raise Exit(1)
 
+    def refresh_credentials(self):
+        with AwsCredsEntryPoint(self, override_entrypoint=""):
+            if exit_code := self._start('echo "Done."'):
+                return exit_code
+
     def start(self, command, *arguments):
         with AwsCredsEntryPoint(self, self.entrypoint):
             return self._start(command, *arguments)
diff --git a/leverage/modules/terraform.py b/leverage/modules/terraform.py
index ce56e70d..04c7c36f 100644
--- a/leverage/modules/terraform.py
+++ b/leverage/modules/terraform.py
@@ -187,6 +187,15 @@ def _import(tf, address, _id):
         raise Exit(exit_code)
 
 
+@terraform.command("refresh-credentials")
+@pass_container
+def refresh_credentials(tf):
+    """Refresh the AWS credentials used on the current layer."""
+    tf.check_for_layer_location()
+    if exit_code := tf.refresh_credentials():
+        raise Exit(exit_code)
+
+
 # ###########################################################################
 # HANDLER FOR MANAGING THE BASE COMMANDS (init, plan, apply, destroy, output)
 # ###########################################################################
diff --git a/tests/test_containers/__init__.py b/tests/test_containers/__init__.py
index f488e09d..14b0844c 100644
--- a/tests/test_containers/__init__.py
+++ b/tests/test_containers/__init__.py
@@ -18,4 +18,5 @@ def container_fixture_factory(container_class):
     with patch("leverage.container.load_env", return_value=FAKE_ENV):
         container = container_class(mocked_client)
         container._run = Mock()
+        container._check_sso_token = Mock()
         return container
diff --git a/tests/test_containers/test_terraform.py b/tests/test_containers/test_terraform.py
index 48d22fe1..7bd8939e 100644
--- a/tests/test_containers/test_terraform.py
+++ b/tests/test_containers/test_terraform.py
@@ -24,3 +24,14 @@ def test_tf_plugin_cache_dir(terraform_container):
 
     # and the cache folder mounted
     assert next(m for m in container_args["host_config"]["Mounts"] if m["Target"] == "/home/testing/.terraform/cache")
+
+
+def test_refresh_credentials(terraform_container):
+    terraform_container.enable_sso()
+    terraform_container.refresh_credentials()
+    container_args = terraform_container.client.api.create_container.call_args_list[0][1]
+
+    # we want a shell, so -> /bin/bash with no entrypoint
+    assert container_args["command"] == 'echo "Done."'
+    # import ipdb; ipdb.set_trace()
+    assert container_args["entrypoint"] == "/root/scripts/aws-sso/aws-sso-entrypoint.sh -- "