Refactor login command

evalai/login.py

@@ -1,9 +1,9 @@
 import os
 import click
-import json
 
 from click import echo, style
 from evalai.utils.auth import get_user_auth_token_by_login
+from evalai.utils.common import store_data_to_json
 from evalai.utils.config import AUTH_TOKEN_PATH, AUTH_TOKEN_DIR
 
 
@@ -17,19 +17,8 @@ def login(ctx):
     password = click.prompt("Enter password", type=str, hide_input=True)
     token = get_user_auth_token_by_login(username, password)
 
-    if os.path.exists(AUTH_TOKEN_PATH):
-        with open(str(AUTH_TOKEN_PATH), "w") as TokenFile:
-            try:
-                json.dump(token, TokenFile)
-            except (OSError, IOError) as e:
-                echo(e)
-    else:
-        if not os.path.exists(AUTH_TOKEN_DIR):
-            os.makedirs(AUTH_TOKEN_DIR)
-        with open(str(AUTH_TOKEN_PATH), "w+") as TokenFile:
-            try:
-                json.dump(token, TokenFile)
-            except (OSError, IOError) as e:
-                echo(e)
+    if not os.path.exists(AUTH_TOKEN_DIR):
+        os.makedirs(AUTH_TOKEN_DIR)
+    store_data_to_json(AUTH_TOKEN_PATH, token, "Unable to store token data due to error: {}")
 
     echo(style("\nLogged in successfully!", bold=True))

evalai/utils/common.py

@@ -1,5 +1,6 @@
-import sys
 import click
+import json
+import sys
 
 from bs4 import BeautifulSoup
 from click import echo, style
@@ -88,3 +89,17 @@ def clean_data(data):
 
 def notify_user(message, color="green", bold=False):
     echo(style(message, fg=color, bold=bold))
+
+
+def store_data_to_json(path, content, message):
+    with open(path, "w") as file:
+        try:
+            json.dump(content, file)
+        except (OSError, IOError) as e:
+            echo(
+                style(
+                    message.format(e),
+                    bold=True,
+                    fg="red",
+                )
+            )

tests/test_login.py

@@ -0,0 +1,35 @@
+import responses
+import json
+
+from unittest.mock import patch
+from click.testing import CliRunner
+from evalai.utils.auth import URLS
+from evalai.utils.config import API_HOST_URL
+from evalai.login import login
+
+from .base import BaseTestClass
+
+
+class TestLogin(BaseTestClass):
+    def setup(self):
+        token = json.loads("""{"token": "test"}""")
+
+        url = "{}{}"
+        responses.add(
+            responses.POST,
+            url.format(API_HOST_URL, URLS.login.value),
+            json=token,
+            status=200,
+        )
+
+    @responses.activate
+    @patch("evalai.utils.common.json.dump")
+    def test_login_when_storing_json_fails(self, mock_dump):
+        error = "Exception"
+        expected = "Unable to store token data due to error: {}".format(error)
+        mock_dump.side_effect = OSError(error)
+
+        runner = CliRunner()
+        result = runner.invoke(login, input="username\npassword")
+        response = result.output
+        assert expected in response