diff --git a/awx/main/management/commands/create_oauth2_token.py b/awx/main/management/commands/create_oauth2_token.py index ca2fd6f05f..2d27566c12 100644 --- a/awx/main/management/commands/create_oauth2_token.py +++ b/awx/main/management/commands/create_oauth2_token.py @@ -1,25 +1,27 @@ -from awx.api.serializers import OAuth2TokenSerializer -from django.core.management.base import BaseCommand +# Django +from django.core.management.base import BaseCommand, CommandError from django.contrib.auth.models import User from django.core.exceptions import ObjectDoesNotExist +# AWX +from awx.api.serializers import OAuth2TokenSerializer + class Command(BaseCommand): """Command that creates an OAuth2 token for a certain user. Returns the value of created token.""" - help='Usage: awx-manage create_oauth2_token --user=username. Will return the value of created token.' + help='Creates an OAuth2 token for a user.' def add_arguments(self, parser): parser.add_argument('--user', dest='user', type=str) def handle(self, *args, **options): if not options['user']: - self.stdout.write("Username not supplied. Usage: awx-manage create_oauth2_token --user=username.") - return + + raise CommandError('Username not supplied. Usage: awx-manage create_oauth2_token --user=username.') try: user = User.objects.get(username=options['user']) except ObjectDoesNotExist: - self.stdout.write("The user does not exist.") - return + raise CommandError('The user does not exist.') config = {'user': user, 'scope':'write'} serializer_obj = OAuth2TokenSerializer() token_record = serializer_obj.create(config, True) diff --git a/awx/main/tests/functional/commands/test_oauth2_token_create.py b/awx/main/tests/functional/commands/test_oauth2_token_create.py index aa15260a07..7b3ceba8e0 100644 --- a/awx/main/tests/functional/commands/test_oauth2_token_create.py +++ b/awx/main/tests/functional/commands/test_oauth2_token_create.py @@ -1,19 +1,27 @@ +# Python import pytest import string import random import StringIO + +# Django from django.contrib.auth.models import User from django.core.management import call_command +from django.core.management.base import CommandError + +# AWX from awx.main.models.oauth import OAuth2AccessToken + @pytest.mark.django_db @pytest.mark.inventory_import class TestOAuth2CreateCommand: def test_no_user_option(self): out = StringIO.StringIO() - call_command('create_oauth2_token', stdout=out) - assert 'Username not supplied.' in out.getvalue() + with pytest.raises(CommandError) as excinfo: + call_command('create_oauth2_token', stdout=out) + assert 'Username not supplied.' in excinfo.value.message out.close() def test_non_existing_user(self): @@ -22,8 +30,9 @@ class TestOAuth2CreateCommand: while fake_username == '' or User.objects.filter(username=fake_username).exists(): fake_username = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(6)) arg = '--user=' + fake_username - call_command('create_oauth2_token', arg, stdout=out) - assert 'The user does not exist.' == out.getvalue().strip() + with pytest.raises(CommandError) as excinfo: + call_command('create_oauth2_token', arg, stdout=out) + assert 'The user does not exist.' in excinfo.value.message out.close() def test_correct_user(self, alice):