facilitate running test environment with Haketilo loaded into browser

1 files changed, 20 insertions, 3 deletions

diff --git a/test/__main__.py b/test/__main__.py
index f8b8976..7afda55 100644
--- a/test/__main__.py
+++ b/test/__main__.py
@@ -36,23 +36,40 @@ import readline
 from .server import do_an_internet
 from .misc_constants import *
 from .profiles import firefox_safe_mode
+from .extension_crafting import get_extension_base_url
 
 def fail(msg, error_code):
     print('Error:', msg)
-    print('Usage:', sys.argv[0], '[certificates_directory] [proxy_port]')
+    print('Usage:', sys.argv[0], '[--load-haketilo]', '[certificates_directory] [proxy_port]')
     sys.exit(error_code)
 
-certdir = Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else default_cert_dir
+load_haketilo = False
+argv_idx = 1
+if len(sys.argv) > argv_idx and sys.argv[argv_idx] == '--load-haketilo':
+    load_haketilo = True
+    argv_idx += 1
+
+certdir = Path(sys.argv[argv_idx]).resolve() if len(sys.argv) > argv_idx \
+    else default_cert_dir
+
 if not certdir.is_dir():
     fail('selected certificate directory does not exist.', 2)
 
-port = sys.argv[2] if len(sys.argv) > 2 else str(default_proxy_port)
+argv_idx += 1
+
+port = sys.argv[argv_idx] if len(sys.argv) > argv_idx \
+    else str(default_proxy_port)
+
 if not port.isnumeric():
     fail('port must be an integer.', 3)
 
 httpd = do_an_internet(certdir, int(port))
 driver = firefox_safe_mode(proxy_port=int(port))
 
+if load_haketilo:
+    driver.install_addon(str(here.parent / 'mozilla-build.zip'), temporary=True)
+    driver.get(get_extension_base_url(driver) + 'html/settings.html')
+
 print("You can now control the browser through 'driver' object")
 
 # Here we enable readline-enhanced editing: