Avoid a backtrace when editing notes on an invalid game

ZLG-Bug: 4

2 files changed, 24 insertions, 1 deletions

diff --git a/src/vgstash_cli.py b/src/vgstash_cli.py
index 323f1d9..d00f339 100755
--- a/src/vgstash_cli.py
+++ b/src/vgstash_cli.py
@@ -162,6 +162,7 @@ def update_game(title, system, attr, val):
         target_game = db.get_game(title, system)
     except:
         click.echo("Game not found. Please try again.")
+        return
     if attr == 'ownership':
         val = vgstash.vtok(val, vgstash.OWNERSHIP)
     if attr == 'progress':
@@ -183,7 +184,12 @@ def update_game(title, system, attr, val):
 @click.option('--edit', '-e', is_flag=True, default=False)
 def notes(title, system, edit):
     db = get_db()
-    target_game = db.get_game(title, system)
+    try:
+        target_game = db.get_game(title, system)
+    except:
+        click.echo("Game not found. Please try again.")
+        return
+
     if edit:
         with tempfile.NamedTemporaryFile() as tmpfile:
             tmpfile.write(target_game.notes.encode("UTF-8"))
diff --git a/tests/test_vgstash_cli.py b/tests/test_vgstash_cli.py
index 30eff62..0e24859 100644
--- a/tests/test_vgstash_cli.py
+++ b/tests/test_vgstash_cli.py
@@ -199,6 +199,14 @@ def test_update():
         'The Legend  |   NES    |   D |   P\n'
     ))
 
+def test_update_invalid():
+    runner = CliRunner()
+    result = runner.invoke(vgstash_cli.cli, ['update', 'Zelda: Skyward Sword', 'NES', 'progress', 'c'])
+    if verbose:
+        print(result.output)
+    assert result.exit_code == 0
+    assert result.output == 'Game not found. Please try again.\n'
+
 
 def test_notes():
     runner = CliRunner()
@@ -214,6 +222,15 @@ def test_notes():
     ))
 
 
+def test_notes_invalid():
+    runner = CliRunner()
+    result = runner.invoke(vgstash_cli.cli, ['notes', 'Zelda: Skyward Sword', 'NES', '-e'])
+    if verbose:
+        print(result.output)
+    assert result.exit_code == 0
+    assert result.output == 'Game not found. Please try again.\n'
+
+
 def test_notes_empty():
     runner = CliRunner()
     result = runner.invoke(vgstash_cli.cli, ['notes', 'Super Mario Bros.', 'NES'])