From 91180b3194d1014d212b64f62996888f15f1ccc5 Mon Sep 17 00:00:00 2001
From: Francisco Rivera <franmrivera@gmail.com>
Date: Sat, 1 Jun 2024 18:48:32 -0300
Subject: [PATCH] tweaks

---
 leverage/container.py | 30 ++++++++++++++----------------
 leverage/path.py      |  5 +++--
 tests/test_path.py    |  2 +-
 3 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/leverage/container.py b/leverage/container.py
index c12c213..63de8c3 100644
--- a/leverage/container.py
+++ b/leverage/container.py
@@ -1,7 +1,6 @@
 import json
 import os
 import re
-import timeit
 import webbrowser
 from io import BytesIO
 from pathlib import Path
@@ -13,15 +12,14 @@
 import dockerpty
 from configupdater import ConfigUpdater
 from docker import DockerClient
-from docker.errors import APIError, NotFound
+from docker.errors import APIError
 from docker.types import Mount
-from typing import Tuple, Union
+from typing import Tuple
 
-from leverage import __toolbox_version__
 from leverage import logger
 from leverage._utils import AwsCredsEntryPoint, CustomEntryPoint, ExitError, ContainerSession
 from leverage.modules.auth import refresh_layer_credentials
-from leverage.logger import console, raw_logger
+from leverage.logger import raw_logger
 from leverage.logger import get_script_log_level
 from leverage.path import PathsHandler
 from leverage.conf import load as load_env
@@ -69,6 +67,7 @@ class LeverageContainer:
 
     LEVERAGE_IMAGE = "binbash/leverage-toolbox"
     SHELL = "/bin/bash"
+    CONTAINER_USER = "leverage"
 
     def __init__(self, client, mounts: tuple = None, env_vars: dict = None):
         """Project related paths are determined and stored. Project configuration is loaded.
@@ -80,7 +79,7 @@ def __init__(self, client, mounts: tuple = None, env_vars: dict = None):
         # Load configs
         self.env_conf = load_env()
 
-        self.paths = PathsHandler(self.env_conf)
+        self.paths = PathsHandler(self.env_conf, self.CONTAINER_USER)
         self.project = self.paths.project
 
         # Set image to use
@@ -141,17 +140,13 @@ def region(self):
         raise ExitError(1, f"No valid region could be found at: {self.paths.cwd.as_posix()}")
 
     @property
-    def local_image(self):
+    def local_image(self) -> BytesIO:
+        """Return the local image that will be built, as a file-like object."""
         return BytesIO(
-            f"""
+            """
             ARG IMAGE_TAG
             FROM binbash/leverage-toolbox:$IMAGE_TAG
-            
-            # Needed as is mounted later on
-            #RUN mkdir /root/.ssh
-            # Needed for git to run propertly
-            #RUN touch /root/.gitconfig
-            
+
             ARG UNAME
             ARG UID
             ARG GID
@@ -164,11 +159,14 @@ def local_image(self):
         )
 
     def ensure_image(self):
+        """
+        Make sure the required local Docker image is available in the system. If not, build it.
+        If the image already exists, re-build it so changes in the arguments can take effect.
+        """
         logger.info(f"Checking for local docker image, tag: {self.image_tag}...")
-        """Make sure the required local Docker image is available in the system. If not, pull it from registry."""
         build_args = {
             "IMAGE_TAG": self.image_tag,
-            "UNAME": "leverage",  # TODO: what is this exactly?
+            "UNAME": self.CONTAINER_USER,
             "GID": str(os.getgid()),
             "UID": str(os.getuid()),
         }
diff --git a/leverage/path.py b/leverage/path.py
index 4e86227..ee622fc 100644
--- a/leverage/path.py
+++ b/leverage/path.py
@@ -134,7 +134,8 @@ class PathsHandler:
     ACCOUNT_TF_VARS = "account.tfvars"
     BACKEND_TF_VARS = "backend.tfvars"
 
-    def __init__(self, env_conf):
+    def __init__(self, env_conf: dict, container_user: str):
+        self.container_user = container_user
         self.home = Path.home()
         self.cwd = Path.cwd()
         try:
@@ -186,7 +187,7 @@ def backend_tfvars(self):
 
     @property
     def guest_aws_credentials_dir(self):
-        return str("/home/leverage/tmp" / Path(self.project))
+        return str(f"/home/{self.container_user}/tmp" / Path(self.project))
 
     @property
     def host_aws_profiles_file(self):
diff --git a/tests/test_path.py b/tests/test_path.py
index c90a4c8..0038fe0 100644
--- a/tests/test_path.py
+++ b/tests/test_path.py
@@ -113,7 +113,7 @@ def test_check_for_cluster_layer(muted_click_context, propagate_logs, caplog):
     """
     Test that if we are not on a cluster layer, we raise an error.
     """
-    paths = PathsHandler({"PROJECT": "test"})
+    paths = PathsHandler({"PROJECT": "test"}, "leverage")
     with patch.object(paths, "check_for_layer_location"):  # assume parent method is already tested
         with pytest.raises(ExitError):
             paths.cwd = Path("/random")