Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions airflow/providers/fab/auth_manager/cli_commands/user_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,26 +85,30 @@ def users_create(args):


def _find_user(args):
with get_application_builder() as appbuilder:
return _find_user_with_appbuilder(args, appbuilder)


def _find_user_with_appbuilder(args, appbuilder):
if not args.username and not args.email:
raise SystemExit("Missing args: must supply one of --username or --email")

if args.username and args.email:
raise SystemExit("Conflicting args: must supply either --username or --email, but not both")

with get_application_builder() as appbuilder:
user = appbuilder.sm.find_user(username=args.username, email=args.email)
if not user:
raise SystemExit(f'User "{args.username or args.email}" does not exist')
user = appbuilder.sm.find_user(username=args.username, email=args.email)
if not user:
raise SystemExit(f'User "{args.username or args.email}" does not exist')
return user


@cli_utils.action_cli
@providers_configuration_loaded
def user_reset_password(args):
"""Reset user password user from DB."""
user = _find_user(args)
password = _create_password(args)
with get_application_builder() as appbuilder:
user = _find_user_with_appbuilder(args, appbuilder)
if appbuilder.sm.reset_password(user.id, password):
print(f'User "{user.username}" password reset successfully')
else:
Expand Down
41 changes: 38 additions & 3 deletions tests/providers/fab/auth_manager/cli_commands/test_user_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import tempfile
from contextlib import redirect_stdout
from io import StringIO
from unittest import mock

import pytest

Expand Down Expand Up @@ -388,9 +389,9 @@ def test_cli_add_user_role(self, create_user_test4):
user_command.users_manage_role(args, remove=False)
assert 'User "test4" added to role "Op"' in stdout.getvalue()

assert _does_user_belong_to_role(
appbuilder=self.appbuilder, email=TEST_USER1_EMAIL, rolename="Op"
), "User should have been added to role 'Op'"
assert _does_user_belong_to_role(appbuilder=self.appbuilder, email=TEST_USER1_EMAIL, rolename="Op"), (
"User should have been added to role 'Op'"
)

def test_cli_remove_user_role(self, create_user_test4):
assert _does_user_belong_to_role(
Expand Down Expand Up @@ -534,3 +535,37 @@ def test_cli_reset_user_password_with_email(self):
with redirect_stdout(StringIO()) as stdout:
user_command.user_reset_password(args)
assert 'User "test3" password reset successfully' in stdout.getvalue()

def test_cli_reset_user_password_uses_single_appbuilder_context(self):
args = self.parser.parse_args(
[
"users",
"create",
"--username",
"test4",
"--lastname",
"doe",
"--firstname",
"jane",
"--email",
"jane@example.com",
"--role",
"Viewer",
"--password",
"TempPass123!",
]
)
user_command.users_create(args)

args = self.parser.parse_args(
["users", "reset-password", "--username", "test4", "--password", "ResetPass456!"]
)

with mock.patch.object(
user_command, "get_application_builder", wraps=user_command.get_application_builder
) as mocked_get_application_builder:
with redirect_stdout(StringIO()) as stdout:
user_command.user_reset_password(args)

assert mocked_get_application_builder.call_count == 1
assert 'User "test4" password reset successfully' in stdout.getvalue()