management: Don't use sys.exit(1).

Using sys.exit in a management command makes it impossible
to unit test the code in question.  The correct approach to do the same
thing in Django management commands is to raise CommandError.

Followup of b570c0dafa
This commit is contained in:
Vishnu Ks
2019-05-04 02:50:39 +05:30
committed by Tim Abbott
parent 687f1bcdaf
commit 123bcea518
26 changed files with 82 additions and 120 deletions

View File

@@ -1,8 +1,7 @@
import sys
from argparse import ArgumentParser
from typing import Any
from django.core.management.base import BaseCommand
from django.core.management.base import BaseCommand, CommandError
from analytics.lib.counts import do_drop_all_analytics_tables
@@ -18,5 +17,4 @@ class Command(BaseCommand):
if options['force']:
do_drop_all_analytics_tables()
else:
print("Would delete all data from analytics tables (!); use --force to do so.")
sys.exit(1)
raise CommandError("Would delete all data from analytics tables (!); use --force to do so.")

View File

@@ -1,8 +1,7 @@
import sys
from argparse import ArgumentParser
from typing import Any
from django.core.management.base import BaseCommand
from django.core.management.base import BaseCommand, CommandError
from analytics.lib.counts import COUNT_STATS, do_drop_single_stat
@@ -20,10 +19,8 @@ class Command(BaseCommand):
def handle(self, *args: Any, **options: Any) -> None:
property = options['property']
if property not in COUNT_STATS:
print("Invalid property: %s" % (property,))
sys.exit(1)
raise CommandError("Invalid property: %s" % (property,))
if not options['force']:
print("No action taken. Use --force.")
sys.exit(1)
raise CommandError("No action taken. Use --force.")
do_drop_single_stat(property)

View File

@@ -2,7 +2,7 @@ import datetime
from argparse import ArgumentParser
from typing import Any, List
from django.core.management.base import BaseCommand
from django.core.management.base import BaseCommand, CommandError
from django.db.models import Count
from django.utils.timezone import now as timezone_now
@@ -73,8 +73,7 @@ class Command(BaseCommand):
try:
realms = [get_realm(string_id) for string_id in options['realms']]
except Realm.DoesNotExist as e:
print(e)
exit(1)
raise CommandError(e)
else:
realms = Realm.objects.all()

View File

@@ -1,7 +1,7 @@
from argparse import ArgumentParser
from typing import Any
from django.core.management.base import BaseCommand
from django.core.management.base import BaseCommand, CommandError
from django.db.models import Q
from zerver.models import Message, Realm, \
@@ -19,8 +19,7 @@ class Command(BaseCommand):
try:
realms = [get_realm(string_id) for string_id in options['realms']]
except Realm.DoesNotExist as e:
print(e)
exit(1)
raise CommandError(e)
else:
realms = Realm.objects.all()

View File

@@ -2,7 +2,7 @@ import datetime
from argparse import ArgumentParser
from typing import Any
from django.core.management.base import BaseCommand
from django.core.management.base import BaseCommand, CommandError
from django.utils.timezone import now as timezone_now
from zerver.models import Message, Realm, Stream, UserProfile, get_realm
@@ -24,8 +24,7 @@ class Command(BaseCommand):
try:
realms = [get_realm(string_id) for string_id in options['realms']]
except Realm.DoesNotExist as e:
print(e)
exit(1)
raise CommandError(e)
else:
realms = Realm.objects.all()

View File

@@ -6,7 +6,7 @@ from typing import Any
import requests
import ujson
from django.conf import settings
from django.core.management.base import BaseCommand
from django.core.management.base import BaseCommand, CommandError
from django.utils.timezone import now as timezone_now
from zerver.models import UserProfile
@@ -33,24 +33,20 @@ class Command(BaseCommand):
if options['api_key'] is None:
try:
if settings.MAILCHIMP_API_KEY is None:
print('MAILCHIMP_API_KEY is None. Check your server settings file.')
exit(1)
raise CommandError('MAILCHIMP_API_KEY is None. Check your server settings file.')
options['api_key'] = settings.MAILCHIMP_API_KEY
except AttributeError:
print('Please supply a MailChimp API key to --api-key, or add a '
'MAILCHIMP_API_KEY to your server settings file.')
exit(1)
raise CommandError('Please supply a MailChimp API key to --api-key, or add a '
'MAILCHIMP_API_KEY to your server settings file.')
if options['list_id'] is None:
try:
if settings.ZULIP_FRIENDS_LIST_ID is None:
print('ZULIP_FRIENDS_LIST_ID is None. Check your server settings file.')
exit(1)
raise CommandError('ZULIP_FRIENDS_LIST_ID is None. Check your server settings file.')
options['list_id'] = settings.ZULIP_FRIENDS_LIST_ID
except AttributeError:
print('Please supply a MailChimp List ID to --list-id, or add a '
'ZULIP_FRIENDS_LIST_ID to your server settings file.')
exit(1)
raise CommandError('Please supply a MailChimp List ID to --list-id, or add a '
'ZULIP_FRIENDS_LIST_ID to your server settings file.')
endpoint = "https://%s.api.mailchimp.com/3.0/lists/%s/members" % \
(options['api_key'].split('-')[1], options['list_id'])

View File

@@ -4,7 +4,7 @@ import time
from typing import Any, Callable, Optional
from django.conf import settings
from django.core.management.base import BaseCommand, CommandParser
from django.core.management.base import BaseCommand, CommandParser, CommandError
from zerver.lib.rate_limiter import RateLimitedUser, \
client, max_api_calls, max_api_window
@@ -44,8 +44,7 @@ than max_api_calls! (trying to trim) %s %s" % (key, count))
def handle(self, *args: Any, **options: Any) -> None:
if not settings.RATE_LIMITING:
print("This machine is not using redis or rate limiting, aborting")
exit(1)
raise CommandError("This machine is not using redis or rate limiting, aborting")
# Find all keys, and make sure they're all within size constraints
wildcard_list = "ratelimit:*:*:list"

View File

@@ -41,8 +41,7 @@ class Command(BaseCommand):
for path in options['gitter_data']:
if not os.path.exists(path):
print("Gitter data file not found: '%s'" % (path,))
exit(1)
raise CommandError("Gitter data file not found: '%s'" % (path,))
# TODO add json check
print("Converting Data ...")
do_convert_data(path, output_dir, num_threads)

View File

@@ -21,7 +21,7 @@ spec:
exporting-from-hipchat-server-or-data-center-for-data-portability-950821555.html
'''
from django.core.management.base import BaseCommand, CommandParser
from django.core.management.base import BaseCommand, CommandParser, CommandError
from zerver.data_import.hipchat import do_convert_data
@@ -56,25 +56,21 @@ class Command(BaseCommand):
output_dir = options["output_dir"]
if output_dir is None:
print("You need to specify --output <output directory>")
exit(1)
raise CommandError("You need to specify --output <output directory>")
if os.path.exists(output_dir) and not os.path.isdir(output_dir):
print(output_dir + " is not a directory")
exit(1)
raise CommandError(output_dir + " is not a directory")
os.makedirs(output_dir, exist_ok=True)
if os.listdir(output_dir):
print('Output directory should be empty!')
exit(1)
raise CommandError('Output directory should be empty!')
output_dir = os.path.realpath(output_dir)
for path in options['hipchat_tar']:
if not os.path.exists(path):
print("Tar file not found: '%s'" % (path,))
exit(1)
raise CommandError("Tar file not found: '%s'" % (path,))
print("Converting Data ...")
do_convert_data(

View File

@@ -40,8 +40,7 @@ class Command(BaseCommand):
token = options['token']
if token is None:
print("Enter slack legacy token!")
exit(1)
raise CommandError("Enter slack legacy token!")
num_threads = int(options['threads'])
if num_threads < 1:
@@ -49,8 +48,7 @@ class Command(BaseCommand):
for path in options['slack_data_zip']:
if not os.path.exists(path):
print("Slack data directory not found: '%s'" % (path,))
exit(1)
raise CommandError("Slack data directory not found: '%s'" % (path,))
print("Converting Data ...")
do_convert_data(path, output_dir, token, threads=num_threads)

View File

@@ -3,7 +3,7 @@ from argparse import ArgumentParser
from typing import Any
from zerver.lib.actions import do_deactivate_user
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.management import ZulipBaseCommand, CommandError
from zerver.lib.sessions import user_sessions
from zerver.models import UserProfile
@@ -39,8 +39,7 @@ class Command(ZulipBaseCommand):
))
if not options["for_real"]:
print("This was a dry run. Pass -f to actually deactivate.")
exit(1)
raise CommandError("This was a dry run. Pass -f to actually deactivate.")
do_deactivate_user(user_profile)
print("Sessions deleted, user deactivated.")

View File

@@ -2,7 +2,7 @@
from argparse import ArgumentParser
from typing import Any
from django.core.management.base import BaseCommand
from django.core.management.base import BaseCommand, CommandError
from zerver.lib.actions import do_delete_old_unclaimed_attachments
from zerver.models import get_old_unclaimed_attachments
@@ -36,8 +36,7 @@ class Command(BaseCommand):
print("")
if not options["for_real"]:
print("This was a dry run. Pass -f to actually delete.")
exit(1)
raise CommandError("This was a dry run. Pass -f to actually delete.")
do_delete_old_unclaimed_attachments(delta_weeks)
print("")

View File

@@ -29,7 +29,7 @@ from imaplib import IMAP4_SSL
from typing import Any, Generator
from django.conf import settings
from django.core.management.base import BaseCommand
from django.core.management.base import BaseCommand, CommandError
from zerver.lib.email_mirror import logger, process_message
@@ -73,8 +73,7 @@ class Command(BaseCommand):
if (not settings.EMAIL_GATEWAY_BOT or not settings.EMAIL_GATEWAY_LOGIN or
not settings.EMAIL_GATEWAY_PASSWORD or not settings.EMAIL_GATEWAY_IMAP_SERVER or
not settings.EMAIL_GATEWAY_IMAP_PORT or not settings.EMAIL_GATEWAY_IMAP_FOLDER):
print("Please configure the Email Mirror Gateway in /etc/zulip/, "
"or specify $ORIGINAL_RECIPIENT if piping a single mail.")
exit(1)
raise CommandError("Please configure the Email Mirror Gateway in /etc/zulip/, "
"or specify $ORIGINAL_RECIPIENT if piping a single mail.")
for message in get_imap_messages():
process_message(message)

View File

@@ -1,14 +1,12 @@
import logging
import sys
from argparse import ArgumentParser
from typing import Any, List, Optional
from django.core.management.base import CommandError
from django.db import connection
from zerver.lib.fix_unreads import fix
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.management import ZulipBaseCommand, CommandError
from zerver.models import Realm, UserProfile
logging.getLogger('zulip.fix_unreads').setLevel(logging.INFO)
@@ -55,8 +53,7 @@ class Command(ZulipBaseCommand):
if options['all']:
if realm is None:
print('You must specify a realm if you choose the --all option.')
sys.exit(1)
raise CommandError('You must specify a realm if you choose the --all option.')
self.fix_all_users(realm)
return

View File

@@ -29,7 +29,7 @@ class Command(ZulipBaseCommand):
if not options['emails']:
self.print_help("./manage.py", "generate_invite_links")
exit(1)
raise CommandError
for email in options['emails']:
try:
@@ -48,10 +48,9 @@ class Command(ZulipBaseCommand):
email_allowed_for_realm(email, realm)
except DomainNotAllowedForRealmError:
if not options["force"]:
print("You've asked to add an external user '%s' to a closed realm '%s'." % (
email, realm.string_id))
print("Are you sure? To do this, pass --force.")
exit(1)
raise CommandError("You've asked to add an external user '{}' to a "
"closed realm '{}'.\nAre you sure? To do this, "
"pass --force.".format(email, realm.string_id))
prereg_user = PreregistrationUser(email=email, realm=realm)
prereg_user.save()

View File

@@ -1,11 +1,10 @@
import sys
from typing import Any
from django.db import ProgrammingError
from confirmation.models import generate_realm_creation_url
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.management import ZulipBaseCommand, CommandError
from zerver.models import Realm
class Command(ZulipBaseCommand):
@@ -22,8 +21,8 @@ class Command(ZulipBaseCommand):
# first check if the db has been initalized
Realm.objects.first()
except ProgrammingError:
print("The Zulip database does not appear to exist. Have you run initialize-database?")
sys.exit(1)
raise CommandError("The Zulip database does not appear to exist. "
"Have you run initialize-database?")
url = generate_realm_creation_url(by_admin=True)
self.stdout.write(self.style.SUCCESS("Please visit the following "

View File

@@ -67,11 +67,10 @@ import a database dump from one or more JSON files."""
for path in options['export_paths']:
path = os.path.realpath(os.path.expanduser(path))
if not os.path.exists(path):
print("Directory not found: '%s'" % (path,))
exit(1)
raise CommandError("Directory not found: '%s'" % (path,))
if not os.path.isdir(path):
print("Export file should be folder; if it's a tarball, please unpack it first.")
exit(1)
raise CommandError("Export file should be folder; if it's a "
"tarball, please unpack it first.")
paths.append(path)
for path in paths:

View File

@@ -8,7 +8,7 @@ from types import FrameType
from typing import Any, List
from django.conf import settings
from django.core.management.base import BaseCommand
from django.core.management.base import BaseCommand, CommandError
from django.utils import autoreload
from zerver.worker.queue_processors import get_active_worker_queues, get_worker
@@ -46,7 +46,7 @@ class Command(BaseCommand):
logger.info("Not using RabbitMQ queue workers in the test suite.")
else:
logger.error("Cannot run a queue processor when USING_RABBITMQ is False!")
sys.exit(1)
raise CommandError
def run_threaded_workers(queues: List[str], logger: logging.Logger) -> None:
cnt = 0

View File

@@ -2,7 +2,7 @@
from argparse import ArgumentParser
from typing import Any
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.management import ZulipBaseCommand, CommandError
from zerver.lib.rate_limiter import RateLimitedUser, \
block_access, unblock_access
from zerver.models import UserProfile, get_user_profile_by_api_key
@@ -38,8 +38,7 @@ class Command(ZulipBaseCommand):
def handle(self, *args: Any, **options: Any) -> None:
if (not options['api_key'] and not options['email']) or \
(options['api_key'] and options['email']):
print("Please enter either an email or API key to manage")
exit(1)
raise CommandError("Please enter either an email or API key to manage")
realm = self.get_realm(options)
if options['email']:
@@ -48,8 +47,7 @@ class Command(ZulipBaseCommand):
try:
user_profile = get_user_profile_by_api_key(options['api_key'])
except UserProfile.DoesNotExist:
print("Unable to get user profile for api key %s" % (options['api_key'],))
exit(1)
raise CommandError("Unable to get user profile for api key %s" % (options['api_key'],))
users = [user_profile]
if options['bots']:

View File

@@ -7,7 +7,7 @@ from django.core.exceptions import ValidationError
from django.db.utils import IntegrityError
from zerver.lib.domains import validate_domain
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.management import ZulipBaseCommand, CommandError
from zerver.models import RealmDomain, get_realm_domains
class Command(ZulipBaseCommand):
@@ -44,23 +44,21 @@ class Command(ZulipBaseCommand):
try:
validate_domain(domain)
except ValidationError as e:
print(e.messages[0])
sys.exit(1)
raise CommandError(e.messages[0])
if options["op"] == "add":
try:
RealmDomain.objects.create(realm=realm, domain=domain,
allow_subdomains=options["allow_subdomains"])
sys.exit(0)
except IntegrityError:
print("The domain %(domain)s is already a part of your organization." % {'domain': domain})
sys.exit(1)
raise CommandError("The domain %(domain)s is already a part "
"of your organization." % {'domain': domain})
elif options["op"] == "remove":
try:
RealmDomain.objects.get(realm=realm, domain=domain).delete()
sys.exit(0)
except RealmDomain.DoesNotExist:
print("No such entry found!")
sys.exit(1)
raise CommandError("No such entry found!")
else:
self.print_help("./manage.py", "realm_domain")
sys.exit(1)
raise CommandError

View File

@@ -4,7 +4,7 @@ from argparse import ArgumentParser
from typing import Any
from zerver.lib.actions import do_add_realm_filter, do_remove_realm_filter
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.management import ZulipBaseCommand, CommandError
from zerver.models import all_realm_filters
class Command(ZulipBaseCommand):
@@ -44,13 +44,13 @@ Example: ./manage.py realm_filters --realm=zulip --op=show
pattern = options['pattern']
if not pattern:
self.print_help("./manage.py", "realm_filters")
sys.exit(1)
raise CommandError
if options["op"] == "add":
url_format_string = options['url_format_string']
if not url_format_string:
self.print_help("./manage.py", "realm_filters")
sys.exit(1)
raise CommandError
do_add_realm_filter(realm, pattern, url_format_string)
sys.exit(0)
elif options["op"] == "remove":
@@ -58,4 +58,4 @@ Example: ./manage.py realm_filters --realm=zulip --op=show
sys.exit(0)
else:
self.print_help("./manage.py", "realm_filters")
sys.exit(1)
raise CommandError

View File

@@ -11,7 +11,7 @@ from django.core.management.base import CommandParser
from zerver.lib.email_mirror import mirror_email_message
from zerver.lib.email_mirror_helpers import encode_email_address
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.management import ZulipBaseCommand, CommandError
from zerver.models import Realm, get_stream, get_realm
@@ -57,7 +57,7 @@ Example:
def handle(self, **options: str) -> None:
if options['fixture'] is None:
self.print_help('./manage.py', 'send_to_email_mirror')
exit(1)
raise CommandError
if options['stream'] is None:
stream = "Denmark"
@@ -93,8 +93,7 @@ Example:
def _parse_email_fixture(self, fixture_path: str) -> Message:
if not self._does_fixture_path_exist(fixture_path):
print('Fixture {} does not exist'.format(fixture_path))
exit(1)
raise CommandError('Fixture {} does not exist'.format(fixture_path))
if fixture_path.endswith('.json'):
message = self._parse_email_json_fixture(fixture_path)

View File

@@ -6,7 +6,7 @@ from django.conf import settings
from django.core.management.base import CommandParser
from django.test import Client
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.management import ZulipBaseCommand, CommandError
from zerver.models import get_realm
class Command(ZulipBaseCommand):
@@ -61,20 +61,18 @@ approach shown above.
headers["HTTP_" + header.upper().replace("-", "_")] = str(custom_headers_dict[header])
return headers
except ValueError as ve:
print('Encountered an error while attempting to parse custom headers: %s' % (ve,))
print('Note: all strings must be enclosed within "" instead of \'\'')
exit(1)
raise CommandError('Encountered an error while attempting to parse custom headers: {}\n'
'Note: all strings must be enclosed within "" instead of \'\''.format(ve))
def handle(self, **options: str) -> None:
if options['fixture'] is None or options['url'] is None:
self.print_help('./manage.py', 'send_webhook_fixture_message')
exit(1)
raise CommandError
full_fixture_path = os.path.join(settings.DEPLOY_ROOT, options['fixture'])
if not self._does_fixture_path_exist(full_fixture_path):
print('Fixture {} does not exist'.format(options['fixture']))
exit(1)
raise CommandError('Fixture {} does not exist'.format(options['fixture']))
headers = self.parse_headers(options['custom-headers'])
json = self._get_fixture_as_json(full_fixture_path)
@@ -90,8 +88,7 @@ approach shown above.
result = client.post(options['url'], json, content_type="application/json",
HTTP_HOST=realm.host)
if result.status_code != 200:
print('Error status %s: %s' % (result.status_code, result.content))
exit(1)
raise CommandError('Error status %s: %s' % (result.status_code, result.content))
def _does_fixture_path_exist(self, fixture_path: str) -> bool:
return os.path.exists(fixture_path)

View File

@@ -7,7 +7,7 @@ from django.core.management.base import CommandParser
from django.db import models
from zerver.lib import utils
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.management import ZulipBaseCommand, CommandError
from zerver.models import UserMessage
class Command(ZulipBaseCommand):
@@ -44,8 +44,7 @@ class Command(ZulipBaseCommand):
def handle(self, *args: Any, **options: Any) -> None:
if not options["flag"] or not options["op"] or not options["email"]:
print("Please specify an operation, a flag and an email")
exit(1)
raise CommandError("Please specify an operation, a flag and an email")
op = options['op']
flag = getattr(UserMessage.flags, options['flag'])
@@ -77,7 +76,7 @@ class Command(ZulipBaseCommand):
if not options["for_real"]:
logging.info("Updating %s by %s %s" % (mids, op, flag))
logging.info("Dry run completed. Run with --for-real to change message flags.")
exit(1)
raise CommandError
utils.run_in_batches(mids, 400, do_update, sleep_time=3)
exit(0)

View File

@@ -63,7 +63,7 @@ class Command(ZulipBaseCommand):
if not user_emails:
print('You need to specify at least one user to use the activate option.')
self.print_help("./manage.py", "soft_deactivate_users")
sys.exit(1)
raise CommandError
users_to_activate = get_users_from_emails(user_emails, filter_kwargs)
users_activated = do_soft_activate_users(users_to_activate)
@@ -81,4 +81,4 @@ class Command(ZulipBaseCommand):
else:
self.print_help("./manage.py", "soft_deactivate_users")
sys.exit(1)
raise CommandError

View File

@@ -169,14 +169,14 @@ class TestSendWebhookFixtureMessage(TestCase):
@patch('zerver.management.commands.send_webhook_fixture_message.Command.print_help')
def test_check_if_command_exits_when_fixture_param_is_empty(self, print_help_mock: MagicMock) -> None:
with self.assertRaises(SystemExit):
with self.assertRaises(CommandError):
call_command(self.COMMAND_NAME, url=self.url)
print_help_mock.assert_any_call('./manage.py', self.COMMAND_NAME)
@patch('zerver.management.commands.send_webhook_fixture_message.Command.print_help')
def test_check_if_command_exits_when_url_param_is_empty(self, print_help_mock: MagicMock) -> None:
with self.assertRaises(SystemExit):
with self.assertRaises(CommandError):
call_command(self.COMMAND_NAME, fixture=self.fixture_path)
print_help_mock.assert_any_call('./manage.py', self.COMMAND_NAME)
@@ -186,7 +186,7 @@ class TestSendWebhookFixtureMessage(TestCase):
self, os_path_exists_mock: MagicMock) -> None:
os_path_exists_mock.return_value = False
with self.assertRaises(SystemExit):
with self.assertRaises(CommandError):
call_command(self.COMMAND_NAME, fixture=self.fixture_path, url=self.url)
os_path_exists_mock.assert_any_call(os.path.join(settings.DEPLOY_ROOT, self.fixture_path))
@@ -206,7 +206,7 @@ class TestSendWebhookFixtureMessage(TestCase):
client = client_mock()
with self.assertRaises(SystemExit):
with self.assertRaises(CommandError):
call_command(self.COMMAND_NAME, fixture=self.fixture_path, url=self.url)
self.assertTrue(ujson_mock.dumps.called)
self.assertTrue(ujson_mock.loads.called)
@@ -225,7 +225,7 @@ class TestSendWebhookFixtureMessage(TestCase):
get_fixture_as_json_mock.return_value = "{}"
improper_headers = '{"X-Custom - Headers": "some_val"}'
with self.assertRaises(SystemExit):
with self.assertRaises(CommandError):
call_command(self.COMMAND_NAME, fixture=self.fixture_path, url=self.url,
custom_headers=improper_headers)
@@ -235,7 +235,7 @@ class TestSendWebhookFixtureMessage(TestCase):
command = Command()
self.assertEqual(command.parse_headers(None), None)
self.assertEqual(command.parse_headers('{"X-Custom-Header": "value"}'), {"HTTP_X_CUSTOM_HEADER": "value"})
with self.assertRaises(SystemExit):
with self.assertRaises(CommandError):
command.parse_headers('{"X-Custom - Headers": "some_val"}')
class TestGenerateRealmCreationLink(ZulipTestCase):