diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index d76e60a11..69f3c22c2 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -24,6 +24,8 @@ def parse_args() -> Namespace: args.no_pr = True args.worktree = True args.effort = "low" + if args.command == "auth": + return args return process_and_validate_cmd_args(args) @@ -370,6 +372,11 @@ def _build_parser() -> ArgumentParser: subparsers.add_parser("vscode-install", help="Install the Codeflash VSCode extension") subparsers.add_parser("init-actions", help="Initialize GitHub Actions workflow") + auth_parser = subparsers.add_parser("auth", help="Authentication commands") + auth_subparsers = auth_parser.add_subparsers(dest="auth_command", help="Auth sub-commands") + auth_subparsers.add_parser("login", help="Log in to Codeflash via OAuth") + auth_subparsers.add_parser("status", help="Check authentication status") + trace_optimize = subparsers.add_parser("optimize", help="Trace and optimize your project.") trace_optimize.add_argument( diff --git a/codeflash/cli_cmds/cmd_auth.py b/codeflash/cli_cmds/cmd_auth.py new file mode 100644 index 000000000..96b863fec --- /dev/null +++ b/codeflash/cli_cmds/cmd_auth.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import os + +import click + +from codeflash.cli_cmds.console import console +from codeflash.cli_cmds.oauth_handler import perform_oauth_signin +from codeflash.code_utils.env_utils import get_codeflash_api_key +from codeflash.code_utils.shell_utils import save_api_key_to_rc +from codeflash.either import is_successful + + +def auth_login() -> None: + """Perform OAuth login and save the API key.""" + try: + existing_api_key = get_codeflash_api_key() + except OSError: + existing_api_key = None + + if existing_api_key: + display_key = f"{existing_api_key[:3]}****{existing_api_key[-4:]}" + console.print(f"[green]Already authenticated with API key {display_key}[/green]") + console.print("To re-authenticate, unset [bold]CODEFLASH_API_KEY[/bold] and run this command again.") + return + + api_key = perform_oauth_signin() + if not api_key: + click.echo("Authentication failed.") + raise SystemExit(1) + + result = save_api_key_to_rc(api_key) + if is_successful(result): + click.echo(result.unwrap()) + else: + click.echo(result.failure()) + + os.environ["CODEFLASH_API_KEY"] = api_key + console.print("[green]Signed in successfully![/green]") + + +def auth_status() -> None: + """Check and display current authentication status.""" + try: + api_key = get_codeflash_api_key() + except OSError: + api_key = None + + if api_key: + display_key = f"{api_key[:3]}****{api_key[-4:]}" + console.print(f"[green]Authenticated[/green] with API key {display_key}") + else: + console.print("[yellow]Not authenticated.[/yellow] Run [bold]codeflash auth login[/bold] to sign in.") diff --git a/codeflash/main.py b/codeflash/main.py index 80d6d156a..93d53a1f3 100644 --- a/codeflash/main.py +++ b/codeflash/main.py @@ -34,7 +34,8 @@ def main() -> None: from codeflash.telemetry.sentry import init_sentry args = parse_args() - print_codeflash_banner() + if args.command != "auth": + print_codeflash_banner() # Check for newer version for all commands check_for_newer_minor_version() @@ -47,7 +48,18 @@ def main() -> None: init_sentry(enabled=not disable_telemetry, exclude_errors=True) posthog_cf.initialize_posthog(enabled=not disable_telemetry) - if args.command == "init": + if args.command == "auth": + from codeflash.cli_cmds.cmd_auth import auth_login, auth_status + + if args.auth_command == "login": + auth_login() + elif args.auth_command == "status": + auth_status() + else: + from codeflash.code_utils.code_utils import exit_with_message + + exit_with_message("Usage: codeflash auth {login,status}", error_on_exit=True) + elif args.command == "init": from codeflash.cli_cmds.cmd_init import init_codeflash init_codeflash() diff --git a/tests/test_cmd_auth.py b/tests/test_cmd_auth.py new file mode 100644 index 000000000..d12cecf58 --- /dev/null +++ b/tests/test_cmd_auth.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from codeflash.cli_cmds.cmd_auth import auth_login +from codeflash.either import Success + + +class TestAuthLogin: + @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key") + @patch("codeflash.cli_cmds.cmd_auth.console") + def test_existing_api_key_skips_oauth(self, mock_console: MagicMock, mock_get_key: MagicMock) -> None: + mock_get_key.return_value = "cf-test1234abcd" + + auth_login() + + mock_console.print.assert_any_call("[green]Already authenticated with API key cf-****abcd[/green]") + mock_console.print.assert_any_call( + "To re-authenticate, unset [bold]CODEFLASH_API_KEY[/bold] and run this command again." + ) + + @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key") + @patch("codeflash.cli_cmds.cmd_auth.console") + def test_existing_api_key_oserror_treated_as_missing( + self, mock_console: MagicMock, mock_get_key: MagicMock + ) -> None: + mock_get_key.side_effect = OSError("permission denied") + + with pytest.raises(SystemExit): + with patch("codeflash.cli_cmds.cmd_auth.perform_oauth_signin", return_value=None): + auth_login() + + @patch("codeflash.cli_cmds.cmd_auth.perform_oauth_signin") + @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key", return_value="") + def test_oauth_failure_exits_with_code_1(self, mock_get_key: MagicMock, mock_oauth: MagicMock) -> None: + mock_oauth.return_value = None + + with pytest.raises(SystemExit, match="1"): + auth_login() + + @patch("codeflash.cli_cmds.cmd_auth.os") + @patch("codeflash.cli_cmds.cmd_auth.save_api_key_to_rc") + @patch("codeflash.cli_cmds.cmd_auth.perform_oauth_signin") + @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key", return_value="") + @patch("codeflash.cli_cmds.cmd_auth.console") + def test_successful_oauth_saves_key( + self, + mock_console: MagicMock, + mock_get_key: MagicMock, + mock_oauth: MagicMock, + mock_save: MagicMock, + mock_os: MagicMock, + ) -> None: + mock_oauth.return_value = "cf-newkey12345678" + mock_save.return_value = Success("API key saved to ~/.zshrc") + + auth_login() + + mock_save.assert_called_once_with("cf-newkey12345678") + mock_os.environ.__setitem__.assert_called_once_with("CODEFLASH_API_KEY", "cf-newkey12345678") + mock_console.print.assert_called_with("[green]Signed in successfully![/green]") + + @patch("codeflash.cli_cmds.cmd_auth.os") + @patch("codeflash.cli_cmds.cmd_auth.save_api_key_to_rc") + @patch("codeflash.cli_cmds.cmd_auth.perform_oauth_signin") + @patch("codeflash.cli_cmds.cmd_auth.get_codeflash_api_key", return_value="") + @patch("codeflash.cli_cmds.cmd_auth.console") + def test_windows_oauth_saves_key( + self, + mock_console: MagicMock, + mock_get_key: MagicMock, + mock_oauth: MagicMock, + mock_save: MagicMock, + mock_os: MagicMock, + ) -> None: + mock_oauth.return_value = "cf-newkey12345678" + mock_os.name = "nt" + mock_save.return_value = Success("API key saved") + + auth_login() + + mock_save.assert_called_once_with("cf-newkey12345678") + mock_os.environ.__setitem__.assert_called_once_with("CODEFLASH_API_KEY", "cf-newkey12345678") + + +class TestAuthSubcommandParsing: + def test_auth_login_parses(self) -> None: + from codeflash.cli_cmds.cli import _build_parser + + _build_parser.cache_clear() + parser = _build_parser() + args = parser.parse_args(["auth", "login"]) + assert args.command == "auth" + assert args.auth_command == "login" + + def test_auth_without_subcommand(self) -> None: + from codeflash.cli_cmds.cli import _build_parser + + _build_parser.cache_clear() + parser = _build_parser() + args = parser.parse_args(["auth"]) + assert args.command == "auth" + assert args.auth_command is None