From a6940092bbf37123e68e2c22418584fa9b4a2d88 Mon Sep 17 00:00:00 2001
From: Cristian <cristian@swapps.com>
Date: Tue, 7 Jul 2020 10:25:02 -0500
Subject: [PATCH] feat: Make sure that depth can only be either 1 or 0

---
 archivebox/cli/archivebox_add.py |  2 +-
 tests/test_args.py               | 10 ++++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/archivebox/cli/archivebox_add.py b/archivebox/cli/archivebox_add.py
index 65335679..2f77f754 100644
--- a/archivebox/cli/archivebox_add.py
+++ b/archivebox/cli/archivebox_add.py
@@ -49,11 +49,11 @@ def main(args: Optional[List[str]]=None, stdin: Optional[IO]=None, pwd: Optional
         "--depth",
         action="store",
         default=0,
+        choices=[0,1],
         type=int,
         help="Recursively archive all linked pages up to this many hops away"
     )
     command = parser.parse_args(args or ())
-    #import_str = accept_stdin(stdin)
     add(
         import_str=command.import_path,
         import_path=None,
diff --git a/tests/test_args.py b/tests/test_args.py
index e0c6020e..91264ef2 100644
--- a/tests/test_args.py
+++ b/tests/test_args.py
@@ -5,7 +5,13 @@ from .fixtures import *
 
 def test_depth_flag_is_accepted(process):
     arg_process = subprocess.run(["archivebox", "add", "https://example.com", "--depth=0"], capture_output=True)
-    assert 'unrecognized arguments: --depth' not in arg_process.stderr.decode('utf-8')
+    assert 'unrecognized arguments: --depth' not in arg_process.stderr.decode("utf-8")
+
+def test_depth_flag_fails_if_it_is_not_0_or_1(process):
+    arg_process = subprocess.run(["archivebox", "add", "https://example.com", "--depth=5"], capture_output=True)
+    assert 'invalid choice' in arg_process.stderr.decode("utf-8")
+    arg_process = subprocess.run(["archivebox", "add", "https://example.com", "--depth=-1"], capture_output=True)
+    assert 'invalid choice' in arg_process.stderr.decode("utf-8")
 
 def test_depth_flag_0_crawls_only_the_arg_page(tmp_path, process):
     arg_process = subprocess.run(["archivebox", "add", "https://example.com", "--depth=0"], capture_output=True)
@@ -19,4 +25,4 @@ def test_depth_flag_1_crawls_the_page_AND_links(tmp_path, process):
     with open(tmp_path / "index.json", "r") as f:
         archive_file = f.read()
     assert "https://example.com" in archive_file
-    assert "https://www.iana.org/domains/example" in archive_file
\ No newline at end of file
+    assert "https://www.iana.org/domains/example" in archive_file