From 08a5059804dd23f16cd67132a6342a0b46752e4a Mon Sep 17 00:00:00 2001
From: Eirini Koutsaniti <eirini.koutsaniti@cscs.ch>
Date: Mon, 17 Jul 2023 14:40:44 +0200
Subject: [PATCH] Use `utilities/whoami` endpoint from FirecREST, when
 possible, to determine the username (#60)

* Add whoami endpoint in pyfirecrest

* Add unittests for whoami endpoint
---
 firecrest/BasicClient.py  | 15 +++++++++++----
 firecrest/cli/__init__.py | 11 ++++++++---
 tests/test_utilities.py   | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/firecrest/BasicClient.py b/firecrest/BasicClient.py
index 17f9e71..fd4ac15 100644
--- a/firecrest/BasicClient.py
+++ b/firecrest/BasicClient.py
@@ -972,13 +972,20 @@ def view(self, machine: str, target_path: str) -> str:
         )
         return self._json_response([resp], 200)["output"]
 
-    def whoami(self) -> Optional[str]:
+    def whoami(self, machine=None) -> Optional[str]:
         """Returns the username that FirecREST will be using to perform the other calls.
-        Will return `None` if the token is not valid.
+        In the case the machine name is passed in the arguments, a call is made to the respective endpoint and the command whoami is run on the machine.
+        Otherwise, the library decodes the token and will return `None` if the token is not valid.
+
+        :calls: GET `/utilities/whoami`
         """
+        if machine:
+            resp = self._get_request(
+                endpoint="/utilities/whoami",
+                additional_headers={"X-Machine-Name": machine},
+            )
+            return self._json_response([resp], 200)["output"]
 
-        # FIXME This needs to be added as an endpoint in FirecREST,
-        # now it's making a guess and it could be wrong.
         try:
             decoded = jwt.decode(
                 self._authorization.get_access_token(),
diff --git a/firecrest/cli/__init__.py b/firecrest/cli/__init__.py
index 4c496d8..bf73e06 100644
--- a/firecrest/cli/__init__.py
+++ b/firecrest/cli/__init__.py
@@ -613,11 +613,16 @@ def tail(
 
 
 @app.command(rich_help_panel="Utilities commands")
-def whoami():
-    """Return the username that FirecREST will be using to perform the other calls
+def whoami(
+    machine: Optional[str] = typer.Argument(
+        None, help="The machine name where the `whoami` command will run."
+    ),
+):
+    """Return the username that FirecREST will be using to perform the other calls.
+    If no machine name is passed the username will be deduced from the token.
     """
     try:
-        console.print(client.whoami())
+        console.print(client.whoami(machine))
     except Exception as e:
         examine_exeption(e)
         raise typer.Exit(code=1)
diff --git a/tests/test_utilities.py b/tests/test_utilities.py
index 4efc545..9a577f7 100644
--- a/tests/test_utilities.py
+++ b/tests/test_utilities.py
@@ -532,6 +532,25 @@ def head_tail_callback(request, uri, response_headers):
     return [status_code, response_headers, json.dumps(ret)]
 
 
+def whoami_callback(request, uri, response_headers):
+    if request.headers["Authorization"] != "Bearer VALID_TOKEN":
+        return [401, response_headers, '{"message": "Bad token; invalid JSON"}']
+
+    if request.headers["X-Machine-Name"] != "cluster1":
+        response_headers["X-Machine-Does-Not-Exist"] = "Machine does not exist"
+        return [
+            400,
+            response_headers,
+            '{"description": "Error on whoami operation", "error": "Machine does not exist"}',
+        ]
+
+    return [
+        200,
+        response_headers,
+        '{"description": "Success on whoami operation.", "output": "username"}'
+    ]
+
+
 @pytest.fixture(autouse=True)
 def setup_callbacks():
     httpretty.enable(allow_net_connect=False, verbose=True)
@@ -612,6 +631,10 @@ def setup_callbacks():
         httpretty.GET, "http://firecrest.cscs.ch/utilities/tail", body=head_tail_callback
     )
 
+    httpretty.register_uri(
+        httpretty.GET, "http://firecrest.cscs.ch/utilities/whoami", body=whoami_callback
+    )
+
     httpretty.register_uri(
         httpretty.POST,
         "https://myauth.com/auth/realms/cscs/protocol/openid-connect/token",
@@ -1299,3 +1322,17 @@ def test_view_invalid_machine(valid_client):
 def test_view_invalid_client(invalid_client):
     with pytest.raises(firecrest.UnauthorizedException):
         invalid_client.view("cluster1", "/path/to/file")
+
+
+def test_whoami(valid_client):
+    assert valid_client.whoami('cluster1') == "username"
+
+
+def test_whoami_invalid_machine(valid_client):
+    with pytest.raises(firecrest.HeaderException):
+        valid_client.whoami('cluster2')
+
+
+def test_whoami_invalid_client(invalid_client):
+    with pytest.raises(firecrest.UnauthorizedException):
+        invalid_client.whoami('cluster1')