# Copyright (c) 2015 Red Hat, Inc.
# 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 collections
from operator import attrgetter

import six
import testscenarios

from neutron import context
from neutron.db import agents_db
from neutron.db import agentschedulers_db
from neutron.db import common_db_mixin
from neutron.scheduler import dhcp_agent_scheduler
from neutron.tests.unit.scheduler import (test_dhcp_agent_scheduler as
                                          test_dhcp_sch)

# Required to generate tests from scenarios. Not compatible with nose.
load_tests = testscenarios.load_tests_apply_scenarios


class BaseTestScheduleNetwork(object):
    """Base class which defines scenarios for schedulers.

        agent_count
            Number of dhcp agents (also number of hosts).

        max_agents_per_network
            Maximum  DHCP Agents that can be scheduled for a network.

        scheduled_agent_count
            Number of agents the network has previously scheduled

        down_agent_count
            Number of dhcp agents which are down

        expected_scheduled_agent_count
            Number of scheduled agents the schedule() should return
            or 'None' if the schedule() cannot schedule the network.
    """

    scenarios = [
        ('No agents scheduled if no agents are present',
         dict(agent_count=0,
              max_agents_per_network=1,
              scheduled_agent_count=0,
              down_agent_count=0,
              expected_scheduled_agent_count=None)),

        ('No agents scheduled if network already hosted and'
         ' max_agents_per_network reached',
         dict(agent_count=1,
              max_agents_per_network=1,
              scheduled_agent_count=1,
              down_agent_count=0,
              expected_scheduled_agent_count=None)),

        ('No agents scheduled if all agents are down',
         dict(agent_count=2,
              max_agents_per_network=1,
              scheduled_agent_count=0,
              down_agent_count=2,
              expected_scheduled_agent_count=None)),

        ('Agent scheduled to the network if network is not yet hosted',
         dict(agent_count=1,
              max_agents_per_network=1,
              scheduled_agent_count=0,
              down_agent_count=0,
              expected_scheduled_agent_count=1)),

        ('Additional Agents scheduled to the network if max_agents_per_network'
         ' is not yet reached',
         dict(agent_count=3,
              max_agents_per_network=3,
              scheduled_agent_count=1,
              down_agent_count=0,
              expected_scheduled_agent_count=2)),

        ('No agent scheduled if agent is dead',
         dict(agent_count=3,
              max_agents_per_network=3,
              scheduled_agent_count=1,
              down_agent_count=1,
              expected_scheduled_agent_count=1)),
    ]


class TestChanceScheduleNetwork(test_dhcp_sch.TestDhcpSchedulerBaseTestCase,
                                agentschedulers_db.DhcpAgentSchedulerDbMixin,
                                agents_db.AgentDbMixin,
                                common_db_mixin.CommonDbMixin,
                                BaseTestScheduleNetwork):
    """Test various scenarios for ChanceScheduler.schedule."""

    def test_schedule_network(self):
        self.config(dhcp_agents_per_network=self.max_agents_per_network)
        scheduler = dhcp_agent_scheduler.ChanceScheduler()

        # create dhcp agents
        hosts = ['host-%s' % i for i in range(self.agent_count)]
        dhcp_agents = self._create_and_set_agents_down(
            hosts, down_agent_count=self.down_agent_count)

        active_agents = dhcp_agents[self.down_agent_count:]

        # schedule some agents before calling schedule
        if self.scheduled_agent_count:
            # schedule the network
            schedule_agents = active_agents[:self.scheduled_agent_count]
            scheduler.resource_filter.bind(self.ctx,
                                           schedule_agents, self.network_id)
        actual_scheduled_agents = scheduler.schedule(self, self.ctx,
                                                     self.network)
        if self.expected_scheduled_agent_count:
            self.assertEqual(self.expected_scheduled_agent_count,
                             len(actual_scheduled_agents))
            hosted_agents = self.list_dhcp_agents_hosting_network(
                self.ctx, self.network_id)
            self.assertEqual(self.scheduled_agent_count +
                             len(actual_scheduled_agents),
                             len(hosted_agents['agents']))
        else:
            self.assertEqual([], actual_scheduled_agents)


class TestWeightScheduleNetwork(test_dhcp_sch.TestDhcpSchedulerBaseTestCase,
                                agentschedulers_db.DhcpAgentSchedulerDbMixin,
                                agents_db.AgentDbMixin,
                                common_db_mixin.CommonDbMixin,
                                BaseTestScheduleNetwork):
    """Test various scenarios for WeightScheduler.schedule."""

    def test_weight_schedule_network(self):
        self.config(dhcp_agents_per_network=self.max_agents_per_network)
        scheduler = dhcp_agent_scheduler.WeightScheduler()

        # create dhcp agents
        hosts = ['host-%s' % i for i in range(self.agent_count)]
        dhcp_agents = self._create_and_set_agents_down(
            hosts, down_agent_count=self.down_agent_count)

        active_agents = dhcp_agents[self.down_agent_count:]

        unscheduled_active_agents = list(active_agents)
        # schedule some agents before calling schedule
        if self.scheduled_agent_count:
            # schedule the network
            schedule_agents = active_agents[:self.scheduled_agent_count]
            scheduler.resource_filter.bind(self.ctx,
                                           schedule_agents, self.network_id)
            for agent in schedule_agents:
                unscheduled_active_agents.remove(agent)
        actual_scheduled_agents = scheduler.schedule(self, self.ctx,
                                                     self.network)
        if self.expected_scheduled_agent_count:
            sorted_unscheduled_active_agents = sorted(
                unscheduled_active_agents,
                key=attrgetter('load'))[0:self.expected_scheduled_agent_count]
            self.assertItemsEqual(
                (agent['id'] for agent in actual_scheduled_agents),
                (agent['id'] for agent in sorted_unscheduled_active_agents))
            self.assertEqual(self.expected_scheduled_agent_count,
                             len(actual_scheduled_agents))
            hosted_agents = self.list_dhcp_agents_hosting_network(
                self.ctx, self.network_id)
            self.assertEqual(self.scheduled_agent_count +
                             len(actual_scheduled_agents),
                             len(hosted_agents['agents']))
        else:
            self.assertEqual([], actual_scheduled_agents)


class TestAutoSchedule(test_dhcp_sch.TestDhcpSchedulerBaseTestCase,
                       agentschedulers_db.DhcpAgentSchedulerDbMixin,
                       agents_db.AgentDbMixin,
                       common_db_mixin.CommonDbMixin):
    """Test various scenarios for ChanceScheduler.auto_schedule_networks.

        Below is the brief description of the scenario variables
        --------------------------------------------------------
        agent_count
            number of DHCP agents (also number of hosts).

        max_agents_per_network
            Maximum  DHCP Agents that can be scheduled for a network.

        network_count
            Number of networks.

        networks_with_dhcp_disabled
            List of networks with dhcp disabled

        hosted_networks
            A mapping of agent id to the ids of the networks that they
            should be initially hosting.

        expected_auto_schedule_return_value
            Expected return value of 'auto_schedule_networks'.

        expected_hosted_networks
            This stores the expected networks that should have been scheduled
            (or that could have already been scheduled) for each agent
            after the 'auto_schedule_networks' function is called.

        no_network_with_az_match
            If this parameter is True, there is no unscheduled network with
            availability_zone_hints matches to an availability_zone of agents
            to be scheduled. The default is False.
    """

    scenarios = [
        ('Agent scheduled to the network if network is not yet hosted',
         dict(agent_count=1,
              max_agents_per_network=1,
              network_count=1,
              networks_with_dhcp_disabled=[],
              hosted_networks={},
              expected_auto_schedule_return_value=True,
              expected_hosted_networks={'agent-0': ['network-0']})),

        ('No agent scheduled if no networks are present',
         dict(agent_count=1,
              max_agents_per_network=1,
              network_count=0,
              networks_with_dhcp_disabled=[],
              hosted_networks={},
              expected_auto_schedule_return_value=False,
              expected_hosted_networks={'agent-0': []})),

        ('Agents scheduled to the networks if networks are not yet hosted',
         dict(agent_count=2,
              max_agents_per_network=3,
              network_count=2,
              networks_with_dhcp_disabled=[],
              hosted_networks={},
              expected_auto_schedule_return_value=True,
              expected_hosted_networks={'agent-0': ['network-0',
                                                    'network-1'],
                                        'agent-1': ['network-0',
                                                    'network-1']})),

        ('No new agents scheduled if networks are already hosted',
         dict(agent_count=2,
              max_agents_per_network=3,
              network_count=2,
              networks_with_dhcp_disabled=[],
              hosted_networks={'agent-0': ['network-0', 'network-1'],
                               'agent-1': ['network-0', 'network-1']},
              expected_auto_schedule_return_value=True,
              expected_hosted_networks={'agent-0': ['network-0',
                                                    'network-1'],
                                        'agent-1': ['network-0',
                                                    'network-1']})),

        ('Additional agents scheduled to the networks if'
         ' max_agents_per_network is not yet reached',
         dict(agent_count=4,
              max_agents_per_network=3,
              network_count=4,
              networks_with_dhcp_disabled=[],
              hosted_networks={'agent-0': ['network-0', 'network-1'],
                               'agent-1': ['network-0'],
                               'agent-2': ['network-2'],
                               'agent-3': ['network-0', 'network-2']},
              expected_auto_schedule_return_value=True,
              expected_hosted_networks={'agent-0': ['network-0',
                                                    'network-1',
                                                    'network-2',
                                                    'network-3'],
                                        'agent-1': ['network-0',
                                                    'network-1',
                                                    'network-2',
                                                    'network-3'],
                                        'agent-2': ['network-1',
                                                    'network-2',
                                                    'network-3'],
                                        'agent-3': ['network-0',
                                                    'network-1',
                                                    'network-2',
                                                    'network-3']})),

        ('No agents scheduled if networks already hosted and'
         ' max_agents_per_network reached',
         dict(agent_count=4,
              max_agents_per_network=1,
              network_count=4,
              networks_with_dhcp_disabled=[],
              hosted_networks={'agent-0': ['network-0'],
                               'agent-1': ['network-2'],
                               'agent-2': ['network-1'],
                               'agent-3': ['network-3']},
              expected_auto_schedule_return_value=True,
              expected_hosted_networks={'agent-0': ['network-0'],
                                        'agent-1': ['network-2'],
                                        'agent-2': ['network-1'],
                                        'agent-3': ['network-3']})),

        ('No agents scheduled to the network with dhcp disabled',
         dict(agent_count=2,
              max_agents_per_network=3,
              network_count=2,
              networks_with_dhcp_disabled=['network-1'],
              hosted_networks={},
              expected_auto_schedule_return_value=True,
              expected_hosted_networks={'agent-0': ['network-0'],
                                        'agent-1': ['network-0']})),

        ('No agents scheduled if all networks have dhcp disabled',
         dict(agent_count=2,
              max_agents_per_network=3,
              network_count=2,
              networks_with_dhcp_disabled=['network-0', 'network-1'],
              hosted_networks={},
              expected_auto_schedule_return_value=False,
              expected_hosted_networks={'agent-0': [],
                                        'agent-1': []})),

        ('No agents scheduled if unscheduled network does not match AZ',
         dict(agent_count=1,
              max_agents_per_network=1,
              network_count=1,
              networks_with_dhcp_disabled=[],
              hosted_networks={},
              expected_auto_schedule_return_value=True,
              expected_hosted_networks={'agent-0': []},
              no_network_with_az_match=True)),
    ]

    def _strip_host_index(self, name):
        """Strips the host index.

        Eg. if name = '2-agent-3', then 'agent-3' is returned.
        """
        return name[name.find('-') + 1:]

    def _extract_index(self, name):
        """Extracts the index number and returns.

        Eg. if name = '2-agent-3', then 3 is returned
        """
        return int(name.split('-')[-1])

    def get_subnets(self, context, fields=None):
        subnets = []
        for net_id in self._networks:
            enable_dhcp = (not self._strip_host_index(net_id) in
                           self.networks_with_dhcp_disabled)
            subnets.append({'network_id': net_id,
                            'enable_dhcp': enable_dhcp})
        return subnets

    def get_network(self, context, net_id):
        az_hints = []
        if getattr(self, 'no_network_with_az_match', False):
            az_hints = ['not-match']
        return {'availability_zone_hints': az_hints}

    def _get_hosted_networks_on_dhcp_agent(self, agent_id):
        query = self.ctx.session.query(
            agentschedulers_db.NetworkDhcpAgentBinding.network_id)
        query = query.filter(
            agentschedulers_db.NetworkDhcpAgentBinding.dhcp_agent_id ==
            agent_id)

        return [item[0] for item in query]

    def _test_auto_schedule(self, host_index):
        self.config(dhcp_agents_per_network=self.max_agents_per_network)
        scheduler = dhcp_agent_scheduler.ChanceScheduler()
        self.ctx = context.get_admin_context()
        msg = 'host_index = %s' % host_index

        # create dhcp agents
        hosts = ['%s-agent-%s' % (host_index, i)
                 for i in range(self.agent_count)]
        dhcp_agents = self._create_and_set_agents_down(hosts)

        # create networks
        self._networks = ['%s-network-%s' % (host_index, i)
                          for i in range(self.network_count)]
        self._save_networks(self._networks)

        # pre schedule the networks to the agents defined in
        # self.hosted_networks before calling auto_schedule_network
        for agent, networks in six.iteritems(self.hosted_networks):
            agent_index = self._extract_index(agent)
            for net in networks:
                net_index = self._extract_index(net)
                scheduler.resource_filter.bind(self.ctx,
                                               [dhcp_agents[agent_index]],
                                               self._networks[net_index])

        retval = scheduler.auto_schedule_networks(self, self.ctx,
                                                  hosts[host_index])
        self.assertEqual(self.expected_auto_schedule_return_value, retval,
                         message=msg)

        agent_id = dhcp_agents[host_index].id
        hosted_networks = self._get_hosted_networks_on_dhcp_agent(agent_id)
        hosted_net_ids = [self._strip_host_index(net)
                          for net in hosted_networks]
        expected_hosted_networks = self.expected_hosted_networks['agent-%s' %
                                                                 host_index]
        self.assertItemsEqual(hosted_net_ids, expected_hosted_networks, msg)

    def test_auto_schedule(self):
        for i in range(self.agent_count):
            self._test_auto_schedule(i)


class TestAZAwareWeightScheduler(test_dhcp_sch.TestDhcpSchedulerBaseTestCase,
                                 agentschedulers_db.DhcpAgentSchedulerDbMixin,
                                 agents_db.AgentDbMixin,
                                 common_db_mixin.CommonDbMixin):
    """Test various scenarios for AZAwareWeightScheduler.schedule.

        az_count
            Number of AZs.

        network_az_hints
            Number of AZs in availability_zone_hints of the network.

        agent_count[each az]
            Number of dhcp agents (also number of hosts).

        max_agents_per_network
            Maximum  DHCP Agents that can be scheduled for a network.

        scheduled_agent_count[each az]
            Number of agents the network has previously scheduled

        down_agent_count[each az]
            Number of dhcp agents which are down

        expected_scheduled_agent_count[each az]
            Number of scheduled agents the schedule() should return
            or 'None' if the schedule() cannot schedule the network.
    """

    scenarios = [
        ('Single hint, Single agent, Scheduled an agent of the specified AZ',
         dict(az_count=2,
              network_az_hints=1,
              agent_count=[1, 1],
              max_agents_per_network=1,
              scheduled_agent_count=[0, 0],
              down_agent_count=[0, 0],
              expected_scheduled_agent_count=[1, 0])),

        ('Multi hints, Multi agents Scheduled agents of the specified AZs',
         dict(az_count=3,
              network_az_hints=2,
              agent_count=[1, 1, 1],
              max_agents_per_network=2,
              scheduled_agent_count=[0, 0, 0],
              down_agent_count=[0, 0, 0],
              expected_scheduled_agent_count=[1, 1, 0])),

        ('Single hint, Multi agents, Scheduled agents of the specified AZ',
         dict(az_count=2,
              network_az_hints=1,
              agent_count=[2, 1],
              max_agents_per_network=2,
              scheduled_agent_count=[0, 0],
              down_agent_count=[0, 0],
              expected_scheduled_agent_count=[2, 0])),

        ('Multi hints, Multi agents, Only single AZ available',
         dict(az_count=2,
              network_az_hints=2,
              agent_count=[2, 1],
              max_agents_per_network=2,
              scheduled_agent_count=[0, 0],
              down_agent_count=[0, 1],
              expected_scheduled_agent_count=[2, 0])),

        ('Multi hints, Multi agents, Not enough agents',
         dict(az_count=3,
              network_az_hints=3,
              agent_count=[1, 1, 1],
              max_agents_per_network=3,
              scheduled_agent_count=[0, 0, 0],
              down_agent_count=[0, 1, 0],
              expected_scheduled_agent_count=[1, 0, 1])),

        ('Multi hints, Multi agents, Partially scheduled, Another AZ selected',
         dict(az_count=3,
              network_az_hints=2,
              agent_count=[1, 1, 1],
              max_agents_per_network=2,
              scheduled_agent_count=[1, 0, 0],
              down_agent_count=[0, 0, 0],
              expected_scheduled_agent_count=[0, 1, 0])),

        ('No hint, Scheduled independent to AZ',
         dict(az_count=3,
              network_az_hints=0,
              agent_count=[1, 1, 1],
              max_agents_per_network=3,
              scheduled_agent_count=[0, 0, 0],
              down_agent_count=[0, 0, 0],
              expected_scheduled_agent_count=[1, 1, 1])),
    ]

    def _set_network_az_hints(self):
        self.network['availability_zone_hints'] = []
        for i in range(self.network_az_hints):
            self.network['availability_zone_hints'].append('az%s' % i)

    def test_schedule_network(self):
        self.config(dhcp_agents_per_network=self.max_agents_per_network)
        scheduler = dhcp_agent_scheduler.AZAwareWeightScheduler()
        self._set_network_az_hints()

        # create dhcp agents
        for i in range(self.az_count):
            az = 'az%s' % i
            hosts = ['%s-host-%s' % (az, j)
                     for j in range(self.agent_count[i])]
            dhcp_agents = self._create_and_set_agents_down(
                hosts, down_agent_count=self.down_agent_count[i], az=az)

            active_agents = dhcp_agents[self.down_agent_count[i]:]

            # schedule some agents before calling schedule
            if self.scheduled_agent_count[i]:
                # schedule the network
                schedule_agents = active_agents[:self.scheduled_agent_count[i]]
                scheduler.resource_filter.bind(
                    self.ctx, schedule_agents, self.network_id)

        actual_scheduled_agents = scheduler.schedule(self, self.ctx,
                                                     self.network)
        scheduled_azs = collections.defaultdict(int)
        for agent in actual_scheduled_agents:
            scheduled_azs[agent['availability_zone']] += 1

        hosted_agents = self.list_dhcp_agents_hosting_network(
                            self.ctx, self.network_id)
        hosted_azs = collections.defaultdict(int)
        for agent in hosted_agents['agents']:
            hosted_azs[agent['availability_zone']] += 1

        for i in range(self.az_count):
            self.assertEqual(self.expected_scheduled_agent_count[i],
                             scheduled_azs.get('az%s' % i, 0))
            self.assertEqual(self.scheduled_agent_count[i] +
                             scheduled_azs.get('az%s' % i, 0),
                             hosted_azs.get('az%s' % i, 0))
