Files
zulip/zerver/management/commands/realm_domain.py
Vishnu Ks 123bcea518 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
2019-05-03 14:20:39 -07:00

65 lines
2.7 KiB
Python

import sys
from argparse import ArgumentParser
from typing import Any
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, CommandError
from zerver.models import RealmDomain, get_realm_domains
class Command(ZulipBaseCommand):
help = """Manage domains for the specified realm"""
def add_arguments(self, parser: ArgumentParser) -> None:
parser.add_argument('--op',
dest='op',
type=str,
default="show",
help='What operation to do (add, show, remove).')
parser.add_argument('--allow-subdomains',
dest='allow_subdomains',
action="store_true",
default=False,
help='Whether subdomains are allowed or not.')
parser.add_argument('domain', metavar='<domain>', type=str, nargs='?',
help="domain to add or remove")
self.add_realm_args(parser, True)
def handle(self, *args: Any, **options: str) -> None:
realm = self.get_realm(options)
assert realm is not None # Should be ensured by parser
if options["op"] == "show":
print("Domains for %s:" % (realm.string_id,))
for realm_domain in get_realm_domains(realm):
if realm_domain["allow_subdomains"]:
print(realm_domain["domain"] + " (subdomains allowed)")
else:
print(realm_domain["domain"] + " (subdomains not allowed)")
sys.exit(0)
domain = options['domain'].strip().lower()
try:
validate_domain(domain)
except ValidationError as e:
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:
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:
raise CommandError("No such entry found!")
else:
self.print_help("./manage.py", "realm_domain")
raise CommandError