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/utils/test_common.py

@@ -0,0 +1,41 @@
+import os
+
+from click.testing import CliRunner
+from evalai.login import login
+from evalai.utils.config import (
+    AUTH_TOKEN_DIR,
+    AUTH_TOKEN_FILE_NAME,
+)
+from unittest.mock import patch
+
+from ..base import BaseTestClass
+
+
+@patch("evalai.login.get_user_auth_token_by_login")
+class TestStoreDataToJson(BaseTestClass):
+    def setup(self):
+        token_file = os.path.join(AUTH_TOKEN_DIR, AUTH_TOKEN_FILE_NAME)
+
+        with open(token_file) as f:
+            self.token = f.read()
+
+    def test_store_data_to_json(self, mock_get_user_auth_token_by_login):
+        mock_get_user_auth_token_by_login.return_value = self.token
+
+        expected = "\n\nLogged in successfully!"
+        runner = CliRunner()
+        result = runner.invoke(login, input="username\npassword")
+        response = result.output
+        assert expected in response
+
+    @patch("evalai.utils.common.json.dump")
+    def test_store_data_to_json_when_error_is_raised(self, mock_dump, mock_get_user_auth_token_by_login):
+        mock_get_user_auth_token_by_login.return_value = self.token
+        error = "Exception"
+        mock_dump.side_effect = OSError(error)
+
+        expected = "Unable to store token data due to error: {}".format(error)
+        runner = CliRunner()
+        result = runner.invoke(login, input="username\npassword")
+        response = result.output
+        assert expected in response