# Copyright 2012 New Dream Network, LLC (DreamHost)
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import copy
import os
import re
import sys
import textwrap
from unittest import mock

from alembic.autogenerate import api as alembic_ag_api
from alembic import config as alembic_config
from alembic.operations import ops as alembic_ops
from alembic import script as alembic_script
import fixtures
from neutron_lib import fixture as lib_fixtures
from neutron_lib.utils import helpers
from oslo_utils import fileutils
import pkg_resources
import sqlalchemy as sa
from testtools import matchers

from neutron.conf.db import migration_cli
from neutron.db import migration
from neutron.db.migration import autogen
from neutron.db.migration import cli
from neutron.tests import base
from neutron.tests import tools
from neutron.tests.unit import testlib_api


class FakeConfig(object):
    service = ''


class FakeRevision(object):
    path = 'fakepath'

    def __init__(self, labels=None, down_revision=None, is_branch_point=False):
        if not labels:
            labels = set()
        self.branch_labels = labels
        self.down_revision = down_revision
        self.is_branch_point = is_branch_point
        self.revision = helpers.get_random_string(10)
        self.module = mock.MagicMock()


class MigrationEntrypointsMemento(fixtures.Fixture):
    '''Create a copy of the migration entrypoints map so it can be restored
       during test cleanup.
    '''

    def _setUp(self):
        self.ep_backup = {}
        for proj, ep in migration_cli.migration_entrypoints.items():
            self.ep_backup[proj] = copy.copy(ep)
        self.addCleanup(self.restore)

    def restore(self):
        migration_cli.migration_entrypoints = self.ep_backup


class TestDbMigration(base.BaseTestCase):

    def setUp(self):
        super(TestDbMigration, self).setUp()
        mock.patch('alembic.op.get_bind').start()
        self.mock_alembic_is_offline = mock.patch(
            'alembic.context.is_offline_mode', return_value=False).start()
        self.mock_alembic_is_o
ffline.return_value = Fals
e
        self.mock_sa_inspector = mock.patch('sqlalchemy.inspect').start()

    def _prepare_mocked_sqlalchemy_inspector(self):
        mock_inspector = mock.MagicMock()
        mock_inspector.get_table_names.return_value = ['foo', 'bar']
        mock_inspector.get_columns.return_value = [{'name': 'foo_column'},
                                                   {'name': 'bar_column'}]
        self.mock_sa_inspector.return_value = mock_inspector

    def test_schema_has_table(self):
        self._prepare_mocked_sqlalchemy_inspector()
        self.assertTrue(migration.schema_has_table('foo'))

    def test_schema_has_table_raises_if_offline(self):
        self.mock_alembic_is_offline.return_value = True
        self.assertRaises(RuntimeError, migration.schema_has_table, 'foo')

    def test_schema_has_column_missing_table(self):
        self._prepare_mocked_sqlalchemy_inspector()
        self.assertFalse(migration.schema_has_column('meh', 'meh'))

    def test_schema_has_column(self):
        self._prepare_mocked_sqlalchemy_inspector()
        self.assertTrue(migration.schema_has_column('foo', 'foo_column'))

    def test_schema_has_column_raises_if_offline(self):
        self.mock_alembic_is_offline.return_value = True
        self.assertRaises(RuntimeError, migration.schema_has_column,
                          'foo', 'foo_col')

    def test_schema_has_column_missing_column(self):
        self._prepare_mocked_sqlalchemy_inspector()
        self.assertFalse(migration.schema_has_column(
            'foo', column_name='meh'))


class TestCli(base.BaseTestCase):
    def setUp(self):
        super(TestCli, self).setUp()
        self.do_alembic_cmd_p = mock.patch.object(cli, 'do_alembic_command')
        self.do_alembic_cmd = self.do_alembic_cmd_p.start()
        self.mock_alembic_err = mock.patch.object(cli, "log_error").start()
        self.mock_alembic_warn = mock.patch.object(cli, "log_warning").start()
        self.mock_alembic_err.side_effect = SystemExit

        def mocked_root_dir(cfg):
            return os.path.join('/fake/dir', cli._get_project_base(cfg))
        mock_root = mock.patch.object(cli, '_get_package_root_dir').start()
        mock_root.side_effect = mocked_root_dir
        # Avoid creating fake directories
        mock.patch('oslo_utils.fileutils.ensure_tree').start()

        # Set up some configs and entrypoints for tests to chew on
        self.configs = []
        self.projects = ('neutron', 'networking-foo', 'neutron-fwaas')
        ini = os.path.join(os.path.dirname(cli.__file__), 'alembic.ini')
        self.useFixture(MigrationEntrypointsMemento())
        migration_cli.migration_entrypoints = {}
        for project in self.projects:
            config = alembic_config.Config(ini)
            config.set_main_option('neutron_project', project)
            module_name = project.replace('-', '_') + '.db.migration'
            attrs = ('alembic_migrations',)
            script_location = ':'.join([module_name, attrs[0]])
            config.set_main_option('script_location', script_location)
            self.configs.append(config)
            entrypoint = pkg_resources.EntryPoint(project,
                                                  module_name,
                                                  attrs=attrs)
            migration_cli.migration_entrypoints[project] = entrypoint

    def _main_test_helper(self, argv, func_name, exp_kwargs=[{}]):
        with mock.patch.object(sys, 'argv', argv),\
                mock.patch.object(cli, 'run_sanity_checks'),\
                mock.patch.object(cli, 'validate_revisions'):

            cli.main()

            def _append_version_path(args):
                args = copy.copy(args)
                if 'autogenerate' in args and not args['autogenerate']:
                    args['version_path'] = mock.ANY
                return args

            self.do_alembic_cmd.assert_has_calls(
                [mock.call(mock.ANY, func_name, **_append_version_path(kwargs))
                 for kwargs in exp_kwargs]
            )

    def test_stamp(self):
        self._main_test_helper(
            ['prog', 'stamp', 'foo'],
            'stamp',
            [{'revision': 'foo', 'sql': False}]
        )

        self._main_test_helper(
            ['prog', 'stamp', 'foo', '--sql'],
            'stamp',
            [{'revision': 'foo', 'sql': True}]
        )

    def _validate_cmd(self, cmd):
        self._main_test_helper(
            ['prog', cmd],
            cmd,
            [{'verbose': False}])

        self._main_test_helper(
            ['prog', cmd, '--verbose'],
            cmd,
            [{'verbose': True}])

    def test_branches(self):
        self._validate_cmd('branches')

    def test_current(self):
        self._validate_cmd('current')

    def test_history(self):
        self._validate_cmd('history')

    def test_heads(self):
        self._validate_cmd('heads')

    def test_check_migration(self):
        with mock.patch.object(cli, 'validate_head_files') as validate:
            self._main_test_helper(['prog', 'check_migration'], 'branches')
            self.assertEqual(len(self.projects), validate.call_count)

    def _test_database_sync_revision(self, separate_branches=True):
        with mock.patch.object(cli, 'update_head_files') as update:
            if separate_branches:
                mock.patch('os.path.exists').start()
            expected_kwargs = [{
                'message': 'message', 'sql': False, 'autogenerate': True,
            }]
   