diff --git a/evalai/login.py b/evalai/login.py index 001f2676c..6064e5591 100644 --- a/evalai/login.py +++ b/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)) diff --git a/evalai/utils/common.py b/evalai/utils/common.py index e59430720..bc4cc4949 100644 --- a/evalai/utils/common.py +++ b/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", + ) + ) diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/utils/test_common.py b/tests/utils/test_common.py new file mode 100644 index 000000000..f8325f2c0 --- /dev/null +++ b/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