Support token override

sqlite_export_for_ynab/_main.py

@@ -50,7 +50,19 @@
 _PACKAGE = "sqlite-export-for-ynab"
 
 
-async def async_main(argv: Sequence[str] | None = None) -> int:
+def resolve_token(token_override: str | None = None) -> str:
+    token = token_override or os.environ.get(_ENV_TOKEN)
+    if token:
+        return token
+
+    raise ValueError(
+        f"Must set YNAB access token as {_ENV_TOKEN!r} environment variable or pass token_override directly. See https://api.ynab.com/#personal-access-tokens"
+    )
+
+
+async def async_main(
+    argv: Sequence[str] | None = None, *, token_override: str | None = None
+) -> int:
     parser = argparse.ArgumentParser(prog=_PACKAGE)
     parser.add_argument(
         "--db",
@@ -71,13 +83,7 @@ async def async_main(argv: Sequence[str] | None = None) -> int:
     db: Path = args.db
     full_refresh: bool = args.full_refresh
 
-    token = os.environ.get(_ENV_TOKEN)
-    if not token:
-        raise ValueError(
-            f"Must set YNAB access token as {_ENV_TOKEN!r} "
-            "environment variable. See "
-            "https://api.ynab.com/#personal-access-tokens"
-        )
+    token = resolve_token(token_override)
 
     await sync(token, db, full_refresh)
 
@@ -530,5 +536,7 @@ async def __call__(
         raise AssertionError("unreachable")
 
 
-def main(argv: Sequence[str] | None = None) -> int:
-    return asyncio.run(async_main(argv))
+def main(
+    argv: Sequence[str] | None = None, *, token_override: str | None = None
+) -> int:
+    return asyncio.run(async_main(argv, token_override=token_override))

tests/_main_test.py

@@ -27,6 +27,7 @@
 from sqlite_export_for_ynab._main import insert_transactions
 from sqlite_export_for_ynab._main import main
 from sqlite_export_for_ynab._main import ProgressYnabClient
+from sqlite_export_for_ynab._main import resolve_token
 from sqlite_export_for_ynab._main import sync
 from sqlite_export_for_ynab._main import YnabClient
 from testing.fixtures import ACCOUNT_ID_1
@@ -586,6 +587,28 @@ def test_main_no_token(tmp_path, monkeypatch):
         main(("--db", str(tmp_path / "db.sqlite")))
 
 
+@patch("sqlite_export_for_ynab._main.sync")
+def test_main_uses_token_override(sync, tmp_path, monkeypatch):
+    monkeypatch.delenv(_ENV_TOKEN, raising=False)
+
+    ret = main(("--db", str(tmp_path / "db.sqlite")), token_override="override-token")
+
+    sync.assert_called_once_with("override-token", tmp_path / "db.sqlite", False)
+    assert ret == 0
+
+
+def test_resolve_token_override(monkeypatch):
+    monkeypatch.delenv(_ENV_TOKEN, raising=False)
+
+    assert resolve_token("override-token") == "override-token"
+
+
+def test_resolve_token_env(monkeypatch):
+    monkeypatch.setenv(_ENV_TOKEN, TOKEN)
+
+    assert resolve_token() == TOKEN
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures(mock_aioresponses.__name__)
 async def test_sync_no_data(tmp_path, mock_aioresponses):