[AIRFLOW-5919] Group tests for the Users/Roles/Perms commands (#6568)

tests/cli/test_cli.py

@@ -33,7 +33,7 @@
 import pytz
 
 import airflow.bin.cli as cli
-from airflow import AirflowException, models, settings
+from airflow import DAG, AirflowException, models, settings
 from airflow.bin.cli import get_dag, get_num_ready_workers_running, run
 from airflow.models import DagModel, Pool, TaskInstance, Variable
 from airflow.settings import Session
@@ -50,6 +50,8 @@
 TEST_DAG_FOLDER = os.path.join(
     os.path.dirname(dag_folder_path), 'dags')
 TEST_DAG_ID = 'unit_tests'
+TEST_USER1_EMAIL = '[email protected]'
+TEST_USER2_EMAIL = '[email protected]'
 
 
 def reset(dag_id):
@@ -533,7 +535,340 @@ def test_dag_state(self):
             'dags', 'state', 'example_bash_operator', DEFAULT_DATE.isoformat()])))
 
 
-class TestCliTest(unittest.TestCase):
+def _does_user_belong_to_role(appbuilder, email, rolename):
+    user = appbuilder.sm.find_user(email=email)
+    role = appbuilder.sm.find_role(rolename)
+    if user and role:
+        return role in user.roles
+
+    return False
+
+
+class TestCliUsers(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.dagbag = models.DagBag(include_examples=True)
+        cls.parser = cli.CLIFactory.get_parser()
+
+    def setUp(self):
+        from airflow.www import app as application
+        self.app, self.appbuilder = application.create_app(session=Session, testing=True)
+        self.clear_roles_and_roles()
+
+    def tearDown(self):
+        self.clear_roles_and_roles()
+
+    def clear_roles_and_roles(self):
+        for email in [TEST_USER1_EMAIL, TEST_USER2_EMAIL]:
+            test_user = self.appbuilder.sm.find_user(email=email)
+            if test_user:
+                self.appbuilder.sm.del_register_user(test_user)
+        for role_name in ['FakeTeamA', 'FakeTeamB']:
+            if self.appbuilder.sm.find_role(role_name):
+                self.appbuilder.sm.delete_role(role_name)
+
+    def test_cli_create_user_random_password(self):
+        args = self.parser.parse_args([
+            'users', 'create', '--username', 'test1', '--lastname', 'doe',
+            '--firstname', 'jon',
+            '--email', '[email protected]', '--role', 'Viewer', '--use_random_password'
+        ])
+        cli.users_create(args)
+
+    def test_cli_create_user_supplied_password(self):
+        args = self.parser.parse_args([
+            'users', 'create', '--username', 'test2', '--lastname', 'doe',
+            '--firstname', 'jon',
+            '--email', '[email protected]', '--role', 'Viewer', '--password', 'test'
+        ])
+        cli.users_create(args)
+
+    def test_cli_delete_user(self):
+        args = self.parser.parse_args([
+            'users', 'create', '--username', 'test3', '--lastname', 'doe',
+            '--firstname', 'jon',
+            '--email', '[email protected]', '--role', 'Viewer', '--use_random_password'
+        ])
+        cli.users_create(args)
+        args = self.parser.parse_args([
+            'users', 'delete', '--username', 'test3',
+        ])
+        cli.users_delete(args)
+
+    def test_cli_list_users(self):
+        for i in range(0, 3):
+            args = self.parser.parse_args([
+                'users', 'create', '--username', 'user{}'.format(i), '--lastname',
+                'doe', '--firstname', 'jon',
+                '--email', 'jdoe+{}@gmail.com'.format(i), '--role', 'Viewer',
+                '--use_random_password'
+            ])
+            cli.users_create(args)
+        with mock.patch('sys.stdout', new_callable=io.StringIO) as mock_stdout:
+            cli.users_list(self.parser.parse_args(['users', 'list']))
+            stdout = mock_stdout.getvalue()
+        for i in range(0, 3):
+            self.assertIn('user{}'.format(i), stdout)
+
+    def test_cli_list_users_with_args(self):
+        cli.users_list(self.parser.parse_args(['users', 'list',
+                                               '--output', 'tsv']))
+
+    def test_cli_import_users(self):
+        def assert_user_in_roles(email, roles):
+            for role in roles:
+                self.assertTrue(_does_user_belong_to_role(self.appbuilder, email, role))
+
+        def assert_user_not_in_roles(email, roles):
+            for role in roles:
+                self.assertFalse(_does_user_belong_to_role(self.appbuilder, email, role))
+
+        assert_user_not_in_roles(TEST_USER1_EMAIL, ['Admin', 'Op'])
+        assert_user_not_in_roles(TEST_USER2_EMAIL, ['Public'])
+        users = [
+            {
+                "username": "imported_user1", "lastname": "doe1",
+                "firstname": "jon", "email": TEST_USER1_EMAIL,
+                "roles": ["Admin", "Op"]
+            },
+            {
+                "username": "imported_user2", "lastname": "doe2",
+                "firstname": "jon", "email": TEST_USER2_EMAIL,
+                "roles": ["Public"]
+            }
+        ]
+        self._import_users_from_file(users)
+
+        assert_user_in_roles(TEST_USER1_EMAIL, ['Admin', 'Op'])
+        assert_user_in_roles(TEST_USER2_EMAIL, ['Public'])
+
+        users = [
+            {
+                "username": "imported_user1", "lastname": "doe1",
+                "firstname": "jon", "email": TEST_USER1_EMAIL,
+                "roles": ["Public"]
+            },
+            {
+                "username": "imported_user2", "lastname": "doe2",
+                "firstname": "jon", "email": TEST_USER2_EMAIL,
+                "roles": ["Admin"]
+            }
+        ]
+        self._import_users_from_file(users)
+
+        assert_user_not_in_roles(TEST_USER1_EMAIL, ['Admin', 'Op'])
+        assert_user_in_roles(TEST_USER1_EMAIL, ['Public'])
+        assert_user_not_in_roles(TEST_USER2_EMAIL, ['Public'])
+        assert_user_in_roles(TEST_USER2_EMAIL, ['Admin'])
+
+    def test_cli_export_users(self):
+        user1 = {"username": "imported_user1", "lastname": "doe1",
+                 "firstname": "jon", "email": TEST_USER1_EMAIL,
+                 "roles": ["Public"]}
+        user2 = {"username": "imported_user2", "lastname": "doe2",
+                 "firstname": "jon", "email": TEST_USER2_EMAIL,
+                 "roles": ["Admin"]}
+        self._import_users_from_file([user1, user2])
+
+        users_filename = self._export_users_to_file()
+        with open(users_filename, mode='r') as file:
+            retrieved_users = json.loads(file.read())
+        os.remove(users_filename)
+
+        # ensure that an export can be imported
+        self._import_users_from_file(retrieved_users)
+
+        def find_by_username(username):
+            matches = [u for u in retrieved_users if u['username'] == username]
+            if not matches:
+                self.fail("Couldn't find user with username {}".format(username))
+                return None
+            else:
+                matches[0].pop('id')  # this key not required for import
+                return matches[0]
+
+        self.assertEqual(find_by_username('imported_user1'), user1)
+        self.assertEqual(find_by_username('imported_user2'), user2)
+
+    def _import_users_from_file(self, user_list):
+        json_file_content = json.dumps(user_list)
+        f = tempfile.NamedTemporaryFile(delete=False)
+        try:
+            f.write(json_file_content.encode())
+            f.flush()
+
+            args = self.parser.parse_args([
+                'users', 'import', f.name
+            ])
+            cli.users_import(args)
+        finally:
+            os.remove(f.name)
+
+    def _export_users_to_file(self):
+        f = tempfile.NamedTemporaryFile(delete=False)
+        args = self.parser.parse_args([
+            'users', 'export', f.name
+        ])
+        cli.users_export(args)
+        return f.name
+
+    def test_cli_add_user_role(self):
+        args = self.parser.parse_args([
+            'users', 'create', '--username', 'test4', '--lastname', 'doe',
+            '--firstname', 'jon',
+            '--email', TEST_USER1_EMAIL, '--role', 'Viewer', '--use_random_password'
+        ])
+        cli.users_create(args)
+
+        self.assertFalse(
+            _does_user_belong_to_role(appbuilder=self.appbuilder, email=TEST_USER1_EMAIL, rolename='Op'),
+            "User should not yet be a member of role 'Op'"
+        )
+
+        args = self.parser.parse_args([
+            'users', 'add_role', '--username', 'test4', '--role', 'Op'
+        ])
+        cli.users_manage_role(args, remove=False)
+
+        self.assertTrue(
+            _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):
+        args = self.parser.parse_args([
+            'users', 'create', '--username', 'test4', '--lastname', 'doe',
+            '--firstname', 'jon',
+            '--email', TEST_USER1_EMAIL, '--role', 'Viewer', '--use_random_password'
+        ])
+        cli.users_create(args)
+
+        self.assertTrue(
+            _does_user_belong_to_role(appbuilder=self.appbuilder, email=TEST_USER1_EMAIL, rolename='Viewer'),
+            "User should have been created with role 'Viewer'"
+        )
+
+        args = self.parser.parse_args([
+            'users', 'remove_role', '--username', 'test4', '--role', 'Viewer'
+        ])
+        cli.users_manage_role(args, remove=True)
+
+        self.assertFalse(
+            _does_user_belong_to_role(appbuilder=self.appbuilder, email=TEST_USER1_EMAIL, rolename='Viewer'),
+            "User should have been removed from role 'Viewer'"
+        )
+
+
+class TestCliSyncPerms(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.dagbag = models.DagBag(include_examples=True)
+        cls.parser = cli.CLIFactory.get_parser()
+
+    def setUp(self):
+        from airflow.www import app as application
+        self.app, self.appbuilder = application.create_app(session=Session, testing=True)
+
+    @mock.patch("airflow.bin.cli.DagBag")
+    def test_cli_sync_perm(self, dagbag_mock):
+        self.expect_dagbag_contains([
+            DAG('has_access_control',
+                access_control={
+                    'Public': {'can_dag_read'}
+                }),
+            DAG('no_access_control')
+        ], dagbag_mock)
+        self.appbuilder.sm = mock.Mock()
+
+        args = self.parser.parse_args([
+            'sync_perm'
+        ])
+        cli.sync_perm(args)
+
+        assert self.appbuilder.sm.sync_roles.call_count == 1
+
+        self.assertEqual(2,
+                         len(self.appbuilder.sm.sync_perm_for_dag.mock_calls))
+        self.appbuilder.sm.sync_perm_for_dag.assert_any_call(
+            'has_access_control',
+            {'Public': {'can_dag_read'}}
+        )
+        self.appbuilder.sm.sync_perm_for_dag.assert_any_call(
+            'no_access_control',
+            None,
+        )
+
+    def expect_dagbag_contains(self, dags, dagbag_mock):
+        dagbag = mock.Mock()
+        dagbag.dags = {dag.dag_id: dag for dag in dags}
+        dagbag_mock.return_value = dagbag
+
+
+class TestCliRoles(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.dagbag = models.DagBag(include_examples=True)
+        cls.parser = cli.CLIFactory.get_parser()
+
+    def setUp(self):
+        from airflow.www import app as application
+        self.app, self.appbuilder = application.create_app(session=Session, testing=True)
+        self.clear_roles_and_roles()
+
+    def tearDown(self):
+        self.clear_roles_and_roles()
+
+    def clear_roles_and_roles(self):
+        for email in [TEST_USER1_EMAIL, TEST_USER2_EMAIL]:
+            test_user = self.appbuilder.sm.find_user(email=email)
+            if test_user:
+                self.appbuilder.sm.del_register_user(test_user)
+        for role_name in ['FakeTeamA', 'FakeTeamB']:
+            if self.appbuilder.sm.find_role(role_name):
+                self.appbuilder.sm.delete_role(role_name)
+
+    def test_cli_create_roles(self):
+        self.assertIsNone(self.appbuilder.sm.find_role('FakeTeamA'))
+        self.assertIsNone(self.appbuilder.sm.find_role('FakeTeamB'))
+
+        args = self.parser.parse_args([
+            'roles', 'create', 'FakeTeamA', 'FakeTeamB'
+        ])
+        cli.roles_create(args)
+
+        self.assertIsNotNone(self.appbuilder.sm.find_role('FakeTeamA'))
+        self.assertIsNotNone(self.appbuilder.sm.find_role('FakeTeamB'))
+
+    def test_cli_create_roles_is_reentrant(self):
+        self.assertIsNone(self.appbuilder.sm.find_role('FakeTeamA'))
+        self.assertIsNone(self.appbuilder.sm.find_role('FakeTeamB'))
+
+        args = self.parser.parse_args([
+            'roles', 'create', 'FakeTeamA', 'FakeTeamB'
+        ])
+
+        cli.roles_create(args)
+
+        self.assertIsNotNone(self.appbuilder.sm.find_role('FakeTeamA'))
+        self.assertIsNotNone(self.appbuilder.sm.find_role('FakeTeamB'))
+
+    def test_cli_list_roles(self):
+        self.appbuilder.sm.add_role('FakeTeamA')
+        self.appbuilder.sm.add_role('FakeTeamB')
+
+        with mock.patch('sys.stdout', new_callable=io.StringIO) as mock_stdout:
+            cli.roles_list(self.parser.parse_args(['roles', 'list']))
+            stdout = mock_stdout.getvalue()
+
+        self.assertIn('FakeTeamA', stdout)
+        self.assertIn('FakeTeamB', stdout)
+
+    def test_cli_list_roles_with_args(self):
+        cli.roles_list(self.parser.parse_args(['roles', 'list',
+                                               '--output', 'tsv']))
+
+
+class TestCliTasks(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         cls.dagbag = models.DagBag(include_examples=True)