Fix credentials generation and credential format control (#9)

* Better control on existing token file

* Better control for credentials formatting

* Remove typo from TFLite Benchmark generation

* Catch unexpected error when reading response from server

edgebenchmark/edgebenchmark.py

@@ -25,6 +25,7 @@
 
 from edgebenchmark.utils import send_model
 from edgebenchmark.utils import load_token_from_file
+from edgebenchmark.utils import CredentialsFormatException
 from edgebenchmark.settings import settings
 from edgebenchmark.custom_types import verify_model_file
 
@@ -40,10 +41,8 @@ def __init__(self, version: int):
 
         try:
             self._token = load_token_from_file()
-        except FileNotFoundError:
-            print(f"{settings._CREDENTIALS_FILE_PATH} file does not exist.\n"
-                  f"Set token with commmand: edgebenchmark configure",
-                  file=sys.stderr)
+        except (FileNotFoundError, CredentialsFormatException):
+            sys.exit(1)
 
     @property
     def version(self):

edgebenchmark/edgebenchmark_cli.py

@@ -24,6 +24,7 @@
 from edgebenchmark.utils import send_model
 from edgebenchmark.utils import get_devices
 from edgebenchmark.utils import load_token_from_file
+from edgebenchmark.utils import CredentialsFormatException
 from edgebenchmark.settings import available_benchmarks
 from edgebenchmark.settings import settings
 from edgebenchmark.custom_types import ModelPathType
@@ -64,12 +65,14 @@ def cli_ncnn():
 
 @cli_configure.command()
 def configure():
+    token_placeholder = "None"
+
     try:
         current_token = load_token_from_file()
-        token_placeholder = current_token[:3] + 3 * "*" + current_token[-3:]
+        if current_token:
+            token_placeholder = current_token[:3] + 3 * "*" + current_token[-3:]
     except FileNotFoundError:
         current_token = ""
-        token_placeholder = "None"
 
     try:
         token = click.prompt(
@@ -79,7 +82,7 @@ def configure():
             default=current_token,
             value_proc=verify_token_size,
         )
-    except ValueError as e:
+    except ValueError:
         print(f"Edge Benchmark Token must have exactly {settings._TOKEN_LENGTH} characters. Please use valid token.")
         sys.exit(1)
 
@@ -90,7 +93,10 @@ def configure():
 
 @cli_devices.command()
 def devices():
-    token = load_token_from_file()
+    try:
+        token = load_token_from_file()
+    except (FileNotFoundError, CredentialsFormatException):
+        sys.exit(1)
 
     response = get_devices(
         settings._PROTOCOL_VERSION,
@@ -438,7 +444,10 @@ def benchmark(
         benchmark_version: str,
         args: Dict[str, Any],
 ):
-    token = load_token_from_file()
+    try:
+        token = load_token_from_file()
+    except (FileNotFoundError, CredentialsFormatException):
+        sys.exit(1)
 
     response = send_model(
         settings._PROTOCOL_VERSION,
@@ -452,9 +461,13 @@ def benchmark(
     )
 
     if response.status_code != 200:
-        response_msg = json.loads(response.content.decode("ascii"))["msg"]
-        print(response_msg, file=sys.stderr)
-        sys.exit(1)
+        try:
+            response_msg = json.loads(response.content.decode("ascii"))["msg"]
+            print(response_msg, file=sys.stderr)
+            sys.exit(1)
+        except Exception:
+            print("Unexpected error", file=sys.stderr)
+            sys.exit(1)
     else:
         print("Model was successfuly sent for benchmarking. Please check the benchmarking result through https://edgebenchmark.com/app website")
 

edgebenchmark/tflite_benchmark.py

@@ -31,7 +31,7 @@ def TFLiteBenchmark(version: int):
     if version in ("1.14.0"):
         b = TFLiteBenchmark_1_14_0
     elif version in ("1.15.0", "1.15.2", "1.15.3", "1.15.4"):
-        b = TFLiteBenchmark_1_15_0(version)
+        b = TFLiteBenchmark_1_15_0
     elif version in ("2.0.0", "2.0.1", "2.0.2", "2.0.3"):
         b = TFLiteBenchmark_2_0_0
     elif version in ("2.1.0", "2.1.1", "2.1.2"):

edgebenchmark/utils.py

@@ -19,13 +19,21 @@
 import hashlib
 from pathlib import Path
 
-from typing import Tuple
-from typing import List
-from typing import Dict
+from typing import (
+    Dict,
+    Optional,
+    List,
+    Tuple,
+)
+
 
 from edgebenchmark.settings import settings
 
 
+class CredentialsFormatException(Exception):
+    pass
+
+
 def send_model(
         protocol_version: Tuple[int, int, int],
         token: str,
@@ -111,7 +119,7 @@ def md5_hash(
     return md5.hexdigest()
 
 
-def load_token_from_file():
+def load_token_from_file() -> Optional[str]:
     settings._CONFIGURE_DIR.mkdir(parents=True, exist_ok=True)
 
     if settings._CREDENTIALS_FILE_PATH.exists():
@@ -122,17 +130,19 @@ def load_token_from_file():
 
                     if key == "edgebenchmark_token":
                         return value
-            except Exception as e:
+            except Exception:
                 print(
                     f"Invalid format of credentials file located at {settings._CONFIGURE_DIR}",
                     file=sys.stderr,
                 )
+                raise CredentialsFormatException
     else:
         print(
             f"{settings._CREDENTIALS_FILE_PATH} file does not exist.\n"
             "Set token with commmand: edgebenchmark configure",
             file=sys.stderr,
         )
+        raise FileNotFoundError
 
 
 def filter_dict(d: Dict):

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="edgebenchmark",
-    version="0.0.10",
+    version="0.0.11",
     author="Bisonai",
     author_email="[email protected]",
     description="Tool to benchmark speed of machine learning models on real mobile devices",