From f9a6dffe1a8433572145bd220381b80623c14d95 Mon Sep 17 00:00:00 2001 From: Vishnu Ks Date: Sun, 20 Aug 2017 00:30:19 +0530 Subject: [PATCH] management: Add add_user_list_args function to ZulipBaseCommand. --- zerver/lib/management.py | 24 ++++++++++++++- zerver/tests/test_management_commands.py | 39 ++++++++++++++++++++---- 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/zerver/lib/management.py b/zerver/lib/management.py index 41e56f934c..41d3181879 100644 --- a/zerver/lib/management.py +++ b/zerver/lib/management.py @@ -5,7 +5,7 @@ from __future__ import print_function from argparse import ArgumentParser from django.core.exceptions import MultipleObjectsReturned from django.core.management.base import BaseCommand, CommandError -from typing import Any, Dict, Optional, Text +from typing import Any, Dict, Optional, Text, List from zerver.models import Realm, UserProfile @@ -31,6 +31,18 @@ You can use the command list_realms to find ID of the realms in this server.""" type=str, help=help) + def add_user_list_args(self, parser, required=False, help=None): + # type: (ArgumentParser, bool, Optional[str]) -> None + if help is None: + help = 'A comma-separated list of email addresses.' + + parser.add_argument( + '-u', '--users', + dest='users', + required=required, + type=str, + help=help) + def get_realm(self, options): # type: (Dict[str, Any]) -> Optional[Realm] val = options["realm_id"] @@ -48,6 +60,16 @@ You can use the command list_realms to find ID of the realms in this server.""" raise CommandError("There is no realm with id '%s'. Aborting." % (options["realm_id"],)) + def get_users(self, options, realm): + # type: (Dict[str, Any], Optional[Realm]) -> List[UserProfile] + if options["users"] is None: + return [] + emails = set([email.strip() for email in options["users"].split(",")]) + user_profiles = [] + for email in emails: + user_profiles.append(self.get_user(email, realm)) + return user_profiles + def get_user(self, email, realm): # type: (Text, Optional[Realm]) -> UserProfile diff --git a/zerver/tests/test_management_commands.py b/zerver/tests/test_management_commands.py index 1cda09b5fb..50861f9962 100644 --- a/zerver/tests/test_management_commands.py +++ b/zerver/tests/test_management_commands.py @@ -17,12 +17,16 @@ from zerver.models import get_realm from confirmation.models import RealmCreationKey, generate_realm_creation_url class TestZulipBaseCommand(ZulipTestCase): + def setUp(self): + # type: () -> None + self.zulip_realm = get_realm("zulip") + def test_get_realm(self): # type: () -> None command = ZulipBaseCommand() - self.assertEqual(command.get_realm(dict(realm_id='zulip')), get_realm("zulip")) + self.assertEqual(command.get_realm(dict(realm_id='zulip')), self.zulip_realm) self.assertEqual(command.get_realm(dict(realm_id=None)), None) - self.assertEqual(command.get_realm(dict(realm_id='1')), get_realm("zulip")) + self.assertEqual(command.get_realm(dict(realm_id='1')), self.zulip_realm) with self.assertRaisesRegex(CommandError, "There is no realm with id"): command.get_realm(dict(realm_id='17')) with self.assertRaisesRegex(CommandError, "There is no realm with id"): @@ -31,19 +35,42 @@ class TestZulipBaseCommand(ZulipTestCase): def test_get_user(self): # type: () -> None command = ZulipBaseCommand() - zulip_realm = get_realm("zulip") mit_realm = get_realm("zephyr") user_profile = self.example_user("hamlet") email = user_profile.email - self.assertEqual(command.get_user(email, zulip_realm), user_profile) + self.assertEqual(command.get_user(email, self.zulip_realm), user_profile) self.assertEqual(command.get_user(email, None), user_profile) with self.assertRaisesRegex(CommandError, "The realm '' does not contain a user with email"): - self.assertEqual(command.get_user(email, mit_realm), user_profile) + command.get_user(email, mit_realm) with self.assertRaisesRegex(CommandError, "server does not contain a user with email"): - self.assertEqual(command.get_user('invalid_email@example.com', None), user_profile) + command.get_user('invalid_email@example.com', None) # TODO: Add a test for the MultipleObjectsReturned case once we make that possible. + def test_get_users(self): + # type: () -> None + command = ZulipBaseCommand() + user_emails = self.example_email("iago") + "," + self.example_email("hamlet") + expected_user_profiles = sorted([self.example_user("iago"), self.example_user("hamlet")], + key = lambda x: x.email) + + self.assertEqual(sorted(command.get_users(dict(users=user_emails), self.zulip_realm), key = lambda x: x.email), + expected_user_profiles) + self.assertEqual(sorted(command.get_users(dict(users=user_emails), None), key = lambda x: x.email), + expected_user_profiles) + + user_emails2 = self.example_email("iago") + "," + self.mit_email("sipbtest") + expected_user_profiles2 = sorted([self.example_user("iago"), self.mit_user("sipbtest")], + key = lambda x: x.email) + self.assertEqual(sorted(command.get_users(dict(users=user_emails2), None), key = lambda x: x.email), + expected_user_profiles2) + with self.assertRaisesRegex(CommandError, "The realm '' does not contain a user with email"): + command.get_users(dict(users=user_emails2), self.zulip_realm) + + self.assertEqual(command.get_users(dict(users=self.example_email("iago")), self.zulip_realm), + [self.example_user("iago")]) + self.assertEqual(command.get_users(dict(users=None), None), []) + class TestCommandsCanStart(TestCase): def setUp(self):