]> review.fuel-infra Code Review - openstack-build/neutron-build.git/commitdiff
Extenuate register_dhcp_agent code duplication in tests
authorAssaf Muller <amuller@redhat.com>
Fri, 1 May 2015 17:29:26 +0000 (13:29 -0400)
committerAssaf Muller <amuller@redhat.com>
Fri, 8 May 2015 19:26:28 +0000 (15:26 -0400)
Non-obvious changes:
* Change helpers.register_agent to use a slimmed down version
  of a plugin that knows how to register an agent. This allows
  the helper to be used with tests that do not register a core
  plugin.

Change-Id: Iefb1af676af6a984b01cdc1e9050541dffb5951a

neutron/tests/common/helpers.py
neutron/tests/functional/scheduler/test_dhcp_agent_scheduler.py
neutron/tests/unit/extensions/test_agent.py
neutron/tests/unit/plugins/openvswitch/test_agent_scheduler.py
neutron/tests/unit/scheduler/test_dhcp_agent_scheduler.py

index bf385117a117cc904a9f9e754cc303ef25644aa7..b780698ad63569ec9cc7c6ba2e2d385a651a03f3 100644 (file)
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
+import datetime
 import os
 
+from oslo_utils import timeutils
+
 import neutron
 from neutron.common import constants
 from neutron.common import topics
 from neutron import context
-from neutron import manager
+from neutron.db import agents_db
+from neutron.db import common_db_mixin
 
 HOST = 'localhost'
 
@@ -37,6 +41,11 @@ def find_sample_file(filename):
         path=os.path.join(neutron.__path__[0], '..', 'etc'))
 
 
+class FakePlugin(common_db_mixin.CommonDbMixin,
+                 agents_db.AgentDbMixin):
+    pass
+
+
 def _get_l3_agent_dict(host, agent_mode, internal_only=True,
                        ext_net_id='', ext_bridge='', router_id=None):
     return {
@@ -53,13 +62,11 @@ def _get_l3_agent_dict(host, agent_mode, internal_only=True,
 
 
 def _register_agent(agent):
-    core_plugin = manager.NeutronManager.get_plugin()
+    plugin = FakePlugin()
     admin_context = context.get_admin_context()
-    core_plugin.create_or_update_agent(admin_context, agent)
-    return core_plugin.get_agents_db(
-        admin_context,
-        filters={'host': [agent['host']],
-                 'agent_type': [agent['agent_type']]})[0]
+    plugin.create_or_update_agent(admin_context, agent)
+    return plugin._get_agent_by_type_and_host(
+        admin_context, agent['agent_type'], agent['host'])
 
 
 def register_l3_agent(host=HOST, agent_mode=constants.L3_AGENT_MODE_LEGACY,
@@ -68,3 +75,46 @@ def register_l3_agent(host=HOST, agent_mode=constants.L3_AGENT_MODE_LEGACY,
     agent = _get_l3_agent_dict(host, agent_mode, internal_only, ext_net_id,
                                ext_bridge, router_id)
     return _register_agent(agent)
+
+
+def _get_dhcp_agent_dict(host, networks=0):
+    agent = {
+        'binary': 'neutron-dhcp-agent',
+        'host': host,
+        'topic': topics.DHCP_AGENT,
+        'agent_type': constants.AGENT_TYPE_DHCP,
+        'configurations': {'dhcp_driver': 'dhcp_driver',
+                           'use_namespaces': True,
+                           'networks': networks}}
+    return agent
+
+
+def register_dhcp_agent(host=HOST, networks=0, admin_state_up=True,
+                        alive=True):
+    agent = _register_agent(
+        _get_dhcp_agent_dict(host, networks))
+
+    if not admin_state_up:
+        set_agent_admin_state(agent['id'])
+    if not alive:
+        kill_agent(agent['id'])
+
+    return FakePlugin()._get_agent_by_type_and_host(
+        context.get_admin_context(), agent['agent_type'], agent['host'])
+
+
+def kill_agent(agent_id):
+    hour_ago = timeutils.utcnow() - datetime.timedelta(hours=1)
+    FakePlugin().update_agent(
+        context.get_admin_context(),
+        agent_id,
+        {'agent': {
+            'started_at': hour_ago,
+            'heartbeat_timestamp': hour_ago}})
+
+
+def set_agent_admin_state(agent_id, admin_state_up=False):
+    FakePlugin().update_agent(
+        context.get_admin_context(),
+        agent_id,
+        {'agent': {'admin_state_up': admin_state_up}})
index 0b0c7d9fdd1977f7d6e35228d06cc18835383ed9..3d37b19411581c9a49c5a9866f95ac9d58dcee9c 100644 (file)
@@ -167,8 +167,9 @@ class TestWeightScheduleNetwork(test_dhcp_sch.TestDhcpSchedulerBaseTestCase,
             sorted_unscheduled_active_agents = sorted(
                 unscheduled_active_agents,
                 key=attrgetter('load'))[0:self.expected_scheduled_agent_count]
-            self.assertItemsEqual(actual_scheduled_agents,
-                                  sorted_unscheduled_active_agents)
+            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(
index 33886b145d1a57c1ca128ce6aae012f332484137..abc28b9aaf45dd2c81ca03b46bc4050434e5f314 100644 (file)
@@ -42,7 +42,6 @@ L3_HOSTA = 'hosta'
 DHCP_HOSTA = 'hosta'
 L3_HOSTB = 'hostb'
 DHCP_HOSTC = 'hostc'
-DHCP_HOST1 = 'host1'
 LBAAS_HOSTA = 'hosta'
 LBAAS_HOSTB = 'hostb'
 
@@ -89,36 +88,24 @@ class AgentDBTestMixIn(object):
             L3_HOSTA, constants.L3_AGENT_MODE_LEGACY)
         l3_hostb = helpers._get_l3_agent_dict(
             L3_HOSTB, constants.L3_AGENT_MODE_LEGACY)
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-        dhcp_hostc = copy.deepcopy(dhcp_hosta)
-        dhcp_hostc['host'] = DHCP_HOSTC
-        lbaas_hosta = {
-            'binary': 'neutron-loadbalancer-agent',
-            'host': LBAAS_HOSTA,
-            'topic': 'LOADBALANCER_AGENT',
-            'configurations': {'device_drivers': ['haproxy_ns']},
-            'agent_type': constants.AGENT_TYPE_LOADBALANCER}
-        lbaas_hostb = copy.deepcopy(lbaas_hosta)
-        lbaas_hostb['host'] = LBAAS_HOSTB
-        callback = agents_db.AgentExtRpcCallback()
+        dhcp_hosta = helpers._get_dhcp_agent_dict(DHCP_HOSTA)
+        dhcp_hostc = helpers._get_dhcp_agent_dict(DHCP_HOSTC)
         helpers.register_l3_agent(host=L3_HOSTA)
         helpers.register_l3_agent(host=L3_HOSTB)
-        callback.report_state(self.adminContext,
-                              agent_state={'agent_state': dhcp_hosta},
-                              time=timeutils.strtime())
-        callback.report_state(self.adminContext,
-                              agent_state={'agent_state': dhcp_hostc},
-                              time=timeutils.strtime())
+        helpers.register_dhcp_agent(host=DHCP_HOSTA)
+        helpers.register_dhcp_agent(host=DHCP_HOSTC)
 
         res = [l3_hosta, l3_hostb, dhcp_hosta, dhcp_hostc]
         if lbaas_agents:
+            lbaas_hosta = {
+                'binary': 'neutron-loadbalancer-agent',
+                'host': LBAAS_HOSTA,
+                'topic': 'LOADBALANCER_AGENT',
+                'configurations': {'device_drivers': ['haproxy_ns']},
+                'agent_type': constants.AGENT_TYPE_LOADBALANCER}
+            lbaas_hostb = copy.deepcopy(lbaas_hosta)
+            lbaas_hostb['host'] = LBAAS_HOSTB
+            callback = agents_db.AgentExtRpcCallback()
             callback.report_state(self.adminContext,
                                   agent_state={'agent_state': lbaas_hosta},
                                   time=timeutils.strtime())
@@ -129,22 +116,6 @@ class AgentDBTestMixIn(object):
 
         return res
 
-    def _register_one_dhcp_agent(self):
-        """Register one DHCP agent."""
-        dhcp_host = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOST1,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-        callback = agents_db.AgentExtRpcCallback()
-        callback.report_state(self.adminContext,
-                              agent_state={'agent_state': dhcp_host},
-                              time=timeutils.strtime())
-        return [dhcp_host]
-
 
 class AgentDBTestCase(AgentDBTestMixIn,
                       test_db_base_plugin_v2.NeutronDbPluginV2TestCase):
@@ -170,13 +141,6 @@ class AgentDBTestCase(AgentDBTestMixIn,
     def test_list_agent(self):
         agents = self._register_agent_states()
         res = self._list('agents')
-        for agt in res['agents']:
-            if (agt['host'] == DHCP_HOSTA and
-                agt['agent_type'] == constants.AGENT_TYPE_DHCP):
-                self.assertEqual(
-                    'dhcp_driver',
-                    agt['configurations']['dhcp_driver'])
-                break
         self.assertEqual(len(agents), len(res['agents']))
 
     def test_show_agent(self):
index 23cfca113face744b7fb19da8fca04f16a2d2273..cd6ea2f7c4913e88c848386d16c33ecc0ee54468 100644 (file)
 # limitations under the License.
 
 import contextlib
-import copy
 import datetime
 
 import mock
 from oslo_config import cfg
 from oslo_db import exception as db_exc
 import oslo_messaging
-from oslo_utils import timeutils
 from webob import exc
 
 from neutron.api import extensions
@@ -182,46 +180,26 @@ class AgentSchedulerTestMixIn(object):
         event_types = [event['event_type'] for event in notifications]
         self.assertIn(expected_event_type, event_types)
 
-    def _register_one_agent_state(self, agent_state):
-        callback = agents_db.AgentExtRpcCallback()
-        callback.report_state(self.adminContext,
-                              agent_state={'agent_state': agent_state},
-                              time=timeutils.strtime())
-
     def test_agent_registration_bad_timestamp(self):
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'start_flag': True,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
         callback = agents_db.AgentExtRpcCallback()
         delta_time = datetime.datetime.now() - datetime.timedelta(days=1)
         str_time = delta_time.strftime('%Y-%m-%dT%H:%M:%S.%f')
-        callback.report_state(self.adminContext,
-                              agent_state={'agent_state': dhcp_hosta},
-                              time=str_time)
+        callback.report_state(
+            self.adminContext,
+            agent_state={
+                'agent_state': helpers._get_dhcp_agent_dict(DHCP_HOSTA)},
+            time=str_time)
 
     def test_agent_registration_invalid_timestamp_allowed(self):
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'start_flag': True,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
         callback = agents_db.AgentExtRpcCallback()
         utc_time = datetime.datetime.utcnow()
         delta_time = utc_time - datetime.timedelta(seconds=10)
         str_time = delta_time.strftime('%Y-%m-%dT%H:%M:%S.%f')
-        callback.report_state(self.adminContext,
-                              agent_state={'agent_state': dhcp_hosta},
-                              time=str_time)
+        callback.report_state(
+            self.adminContext,
+            agent_state={
+                'agent_state': helpers._get_dhcp_agent_dict(DHCP_HOSTA)},
+            time=str_time)
 
     def _disable_agent(self, agent_id, admin_state_up=False):
         new_agent = {}
@@ -393,27 +371,17 @@ class OvsAgentSchedulerTestCase(OvsAgentSchedulerTestCaseBase):
     def test_network_auto_schedule_with_hosted_2(self):
         # one agent hosts one network
         dhcp_rpc_cb = dhcp_rpc.DhcpRpcCallback()
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-        dhcp_hostc = copy.deepcopy(dhcp_hosta)
-        dhcp_hostc['host'] = DHCP_HOSTC
         cfg.CONF.set_override('allow_overlapping_ips', True)
         with self.subnet() as sub1:
-            self._register_one_agent_state(dhcp_hosta)
+            helpers.register_dhcp_agent(DHCP_HOSTA)
             dhcp_rpc_cb.get_active_networks(self.adminContext, host=DHCP_HOSTA)
             hosta_id = self._get_agent_id(constants.AGENT_TYPE_DHCP,
                                           DHCP_HOSTA)
             self._disable_agent(hosta_id, admin_state_up=False)
             with self.subnet() as sub2:
-                self._register_one_agent_state(dhcp_hostc)
+                helpers.register_dhcp_agent(DHCP_HOSTC)
                 dhcp_rpc_cb.get_active_networks(self.adminContext,
-                                             host=DHCP_HOSTC)
+                                                host=DHCP_HOSTC)
                 dhcp_agents_1 = self._list_dhcp_agents_hosting_network(
                     sub1['subnet']['network_id'])
                 dhcp_agents_2 = self._list_dhcp_agents_hosting_network(
@@ -474,7 +442,7 @@ class OvsAgentSchedulerTestCase(OvsAgentSchedulerTestCaseBase):
                 dhcp_agents = self._list_dhcp_agents_hosting_network(
                     port['port']['network_id'])
                 result1 = len(dhcp_agents['agents'])
-            self._register_one_dhcp_agent()
+            helpers.register_dhcp_agent('host1')
             with self.port(subnet=subnet,
                            device_owner="compute:test:" + DHCP_HOSTA) as port:
                 dhcp_agents = self._list_dhcp_agents_hosting_network(
@@ -485,15 +453,7 @@ class OvsAgentSchedulerTestCase(OvsAgentSchedulerTestCaseBase):
         self.assertEqual(3, result2)
 
     def test_network_scheduler_with_disabled_agent(self):
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-        self._register_one_agent_state(dhcp_hosta)
+        helpers.register_dhcp_agent(DHCP_HOSTA)
         with self.port() as port1:
             dhcp_agents = self._list_dhcp_agents_hosting_network(
                 port1['port']['network_id'])
@@ -530,15 +490,7 @@ class OvsAgentSchedulerTestCase(OvsAgentSchedulerTestCaseBase):
                                                               None, None))
 
     def test_network_scheduler_with_down_agent(self):
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-        self._register_one_agent_state(dhcp_hosta)
+        helpers.register_dhcp_agent(DHCP_HOSTA)
         eligible_agent_str = ('neutron.db.agentschedulers_db.'
                               'DhcpAgentSchedulerDbMixin.is_eligible_agent')
         with mock.patch(eligible_agent_str) as eligible_agent:
@@ -560,15 +512,7 @@ class OvsAgentSchedulerTestCase(OvsAgentSchedulerTestCaseBase):
 
     def test_network_scheduler_with_hosted_network(self):
         plugin = manager.NeutronManager.get_plugin()
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-        self._register_one_agent_state(dhcp_hosta)
+        helpers.register_dhcp_agent(DHCP_HOSTA)
         with self.port() as port1:
             dhcp_agents = self._list_dhcp_agents_hosting_network(
                 port1['port']['network_id'])
@@ -640,17 +584,8 @@ class OvsAgentSchedulerTestCase(OvsAgentSchedulerTestCaseBase):
         self._test_network_add_to_dhcp_agent(admin_state_up=False)
 
     def test_network_remove_from_dhcp_agent(self):
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-        self._register_one_agent_state(dhcp_hosta)
-        hosta_id = self._get_agent_id(constants.AGENT_TYPE_DHCP,
-                                      DHCP_HOSTA)
+        agent = helpers.register_dhcp_agent(DHCP_HOSTA)
+        hosta_id = agent.id
         with self.port() as port1:
             num_before_remove = len(
                 self._list_networks_hosted_by_dhcp_agent(
@@ -670,15 +605,7 @@ class OvsAgentSchedulerTestCase(OvsAgentSchedulerTestCaseBase):
         self.assertEqual([], nets)
 
     def test_reserved_port_after_network_remove_from_dhcp_agent(self):
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-        self._register_one_agent_state(dhcp_hosta)
+        helpers.register_dhcp_agent(DHCP_HOSTA)
         hosta_id = self._get_agent_id(constants.AGENT_TYPE_DHCP,
                                       DHCP_HOSTA)
         with self.port(device_owner=constants.DEVICE_OWNER_DHCP,
@@ -698,15 +625,7 @@ class OvsAgentSchedulerTestCase(OvsAgentSchedulerTestCaseBase):
         if keep_services:
             cfg.CONF.set_override(
                 'enable_services_on_agents_with_admin_state_down', True)
-        dhcp_hosta = {
-            'binary': 'neutron-dhcp-agent',
-            'host': DHCP_HOSTA,
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-        self._register_one_agent_state(dhcp_hosta)
+        helpers.register_dhcp_agent(DHCP_HOSTA)
         dhcp_rpc_cb = dhcp_rpc.DhcpRpcCallback()
         with self.port():
             nets = dhcp_rpc_cb.get_active_networks(self.adminContext,
@@ -1312,15 +1231,9 @@ class OvsDhcpAgentNotifierTestCase(test_l3.L3NatTestCaseMixin,
                 {'admin_state_up': False}, DHCP_HOSTA)
 
     def _network_port_create(
-        self, hosts, gateway=attributes.ATTR_NOT_SPECIFIED, owner=None):
+            self, hosts, gateway=attributes.ATTR_NOT_SPECIFIED, owner=None):
         for host in hosts:
-            self._register_one_agent_state(
-                {'binary': 'neutron-dhcp-agent',
-                 'host': host,
-                 'topic': 'dhcp_agent',
-                 'configurations': {'dhcp_driver': 'dhcp_driver',
-                                    'use_namespaces': True, },
-                 'agent_type': constants.AGENT_TYPE_DHCP})
+            helpers.register_dhcp_agent(host)
         with self.network() as net1:
             with self.subnet(network=net1,
                              gateway_ip=gateway) as subnet1:
@@ -1493,12 +1406,10 @@ class OvsL3AgentNotifierTestCase(test_l3.L3NatTestCaseMixin,
         ) as (
             mock_prepare, mock_cast
         ):
-            self._register_agent_states()
-            hosta_id = self._get_agent_id(constants.AGENT_TYPE_L3,
-                                          L3_HOSTA)
-            self._disable_agent(hosta_id, admin_state_up=False)
+            agent_id = helpers.register_l3_agent(L3_HOSTA).id
+            self._disable_agent(agent_id, admin_state_up=False)
 
-            mock_prepare.assert_called_with(server='hosta')
+            mock_prepare.assert_called_with(server=L3_HOSTA)
 
             mock_cast.assert_called_with(
                 mock.ANY, 'agent_updated', payload={'admin_state_up': False})
index 44b5850eddeff47b2d1d4a85698b48d215c0d416..c43e2a22055551ac12517350c38f12aaac601ffc 100644 (file)
 # limitations under the License.
 
 import contextlib
-import datetime
 
 import mock
 from oslo_config import cfg
 from oslo_utils import importutils
-from oslo_utils import timeutils
 import testscenarios
 
 from neutron.common import constants
-from neutron.common import topics
 from neutron import context
-from neutron.db import agents_db
 from neutron.db import agentschedulers_db as sched_db
 from neutron.db import models_v2
 from neutron.extensions import dhcpagentscheduler
 from neutron.scheduler import dhcp_agent_scheduler
+from neutron.tests.common import helpers
 from neutron.tests.unit import testlib_api
 
 # Required to generate tests from scenarios. Not compatible with nose.
 load_tests = testscenarios.load_tests_apply_scenarios
 
+HOST_C = 'host-c'
+HOST_D = 'host-d'
+
 
 class TestDhcpSchedulerBaseTestCase(testlib_api.SqlTestCase):
 
@@ -45,37 +45,16 @@ class TestDhcpSchedulerBaseTestCase(testlib_api.SqlTestCase):
         self.network_id = 'foo_network_id'
         self._save_networks([self.network_id])
 
-    def _get_agents(self, hosts):
-        return [
-            agents_db.Agent(
-                binary='neutron-dhcp-agent',
-                host=host,
-                topic=topics.DHCP_AGENT,
-                configurations="",
-                agent_type=constants.AGENT_TYPE_DHCP,
-                created_at=timeutils.utcnow(),
-                started_at=timeutils.utcnow(),
-                heartbeat_timestamp=timeutils.utcnow())
-            for host in hosts
-        ]
-
-    def _save_agents(self, agents):
-        for agent in agents:
-            with self.ctx.session.begin(subtransactions=True):
-                self.ctx.session.add(agent)
-
-    def _create_and_set_agents_down(self, hosts, down_agent_count=0, **kwargs):
-        dhcp_agents = self._get_agents(hosts)
-        # bring down the specified agents
-        for agent in dhcp_agents[:down_agent_count]:
-            old_time = agent['heartbeat_timestamp']
-            hour_old = old_time - datetime.timedelta(hours=1)
-            agent['heartbeat_timestamp'] = hour_old
-            agent['started_at'] = hour_old
-        for agent in dhcp_agents:
-            agent.update(kwargs)
-        self._save_agents(dhcp_agents)
-        return dhcp_agents
+    def _create_and_set_agents_down(self, hosts, down_agent_count=0,
+                                    admin_state_up=True):
+        agents = []
+        for i, host in enumerate(hosts):
+            is_alive = i >= down_agent_count
+            agents.append(helpers.register_dhcp_agent(
+                host,
+                admin_state_up=admin_state_up,
+                alive=is_alive))
+        return agents
 
     def _save_networks(self, networks):
         for network_id in networks:
@@ -273,26 +252,6 @@ class TestNetworksFailover(TestDhcpSchedulerBaseTestCase,
 class DHCPAgentWeightSchedulerTestCase(TestDhcpSchedulerBaseTestCase):
     """Unit test scenarios for WeightScheduler.schedule."""
 
-    hostc = {
-            'binary': 'neutron-dhcp-agent',
-            'host': 'host-c',
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'networks': 0,
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-
-    hostd = {
-            'binary': 'neutron-dhcp-agent',
-            'host': 'host-d',
-            'topic': 'DHCP_AGENT',
-            'configurations': {'dhcp_driver': 'dhcp_driver',
-                               'networks': 1,
-                               'use_namespaces': True,
-                               },
-            'agent_type': constants.AGENT_TYPE_DHCP}
-
     def setUp(self):
         super(DHCPAgentWeightSchedulerTestCase, self).setUp()
         DB_PLUGIN_KLASS = 'neutron.plugins.ml2.plugin.Ml2Plugin'
@@ -313,10 +272,9 @@ class DHCPAgentWeightSchedulerTestCase(TestDhcpSchedulerBaseTestCase):
         cfg.CONF.set_override("dhcp_load_type", "networks")
 
     def test_scheduler_one_agents_per_network(self):
-        cfg.CONF.set_override('dhcp_agents_per_network', 1)
         self._save_networks(['1111'])
-        agents = self._get_agents(['host-c', 'host-d'])
-        self._save_agents(agents)
+        helpers.register_dhcp_agent(HOST_C)
+        helpers.register_dhcp_agent(HOST_C)
         self.plugin.network_scheduler.schedule(self.plugin, self.ctx,
                                                {'id': '1111'})
         agents = self.plugin.get_dhcp_agents_hosting_networks(self.ctx,
@@ -326,8 +284,8 @@ class DHCPAgentWeightSchedulerTestCase(TestDhcpSchedulerBaseTestCase):
     def test_scheduler_two_agents_per_network(self):
         cfg.CONF.set_override('dhcp_agents_per_network', 2)
         self._save_networks(['1111'])
-        agents = self._get_agents(['host-c', 'host-d'])
-        self._save_agents(agents)
+        helpers.register_dhcp_agent(HOST_C)
+        helpers.register_dhcp_agent(HOST_D)
         self.plugin.network_scheduler.schedule(self.plugin, self.ctx,
                                                {'id': '1111'})
         agents = self.plugin.get_dhcp_agents_hosting_networks(self.ctx,
@@ -343,37 +301,23 @@ class DHCPAgentWeightSchedulerTestCase(TestDhcpSchedulerBaseTestCase):
         self.assertEqual(0, len(agents))
 
     def test_scheduler_equal_distribution(self):
-        cfg.CONF.set_override('dhcp_agents_per_network', 1)
         self._save_networks(['1111', '2222', '3333'])
-        agents = self._get_agents(['host-c', 'host-d'])
-        self._save_agents(agents)
-        callback = agents_db.AgentExtRpcCallback()
-        callback.report_state(self.ctx,
-                              agent_state={'agent_state': self.hostc},
-                              time=timeutils.strtime())
-        callback.report_state(self.ctx,
-                              agent_state={'agent_state': self.hostd},
-                              time=timeutils.strtime())
-        self.plugin.network_scheduler.schedule(self.plugin, self.ctx,
-                                               {'id': '1111'})
-        agent1 = self.plugin.get_dhcp_agents_hosting_networks(self.ctx,
-                                                              ['1111'])
-        self.hostd['configurations']['networks'] = 2
-        callback.report_state(self.ctx,
-                              agent_state={'agent_state': self.hostd},
-                              time=timeutils.strtime())
-        self.plugin.network_scheduler.schedule(self.plugin, self.ctx,
-                                               {'id': '2222'})
-        agent2 = self.plugin.get_dhcp_agents_hosting_networks(self.ctx,
-                                                              ['2222'])
-        self.hostc['configurations']['networks'] = 4
-        callback.report_state(self.ctx,
-                              agent_state={'agent_state': self.hostc},
-                              time=timeutils.strtime())
-        self.plugin.network_scheduler.schedule(self.plugin, self.ctx,
-                                               {'id': '3333'})
-        agent3 = self.plugin.get_dhcp_agents_hosting_networks(self.ctx,
-                                                              ['3333'])
+        helpers.register_dhcp_agent(HOST_C)
+        helpers.register_dhcp_agent(HOST_D, networks=1)
+        self.plugin.network_scheduler.schedule(
+            self.plugin, context.get_admin_context(), {'id': '1111'})
+        helpers.register_dhcp_agent(HOST_D, networks=2)
+        self.plugin.network_scheduler.schedule(
+            self.plugin, context.get_admin_context(), {'id': '2222'})
+        helpers.register_dhcp_agent(HOST_C, networks=4)
+        self.plugin.network_scheduler.schedule(
+            self.plugin, context.get_admin_context(), {'id': '3333'})
+        agent1 = self.plugin.get_dhcp_agents_hosting_networks(
+            self.ctx, ['1111'])
+        agent2 = self.plugin.get_dhcp_agents_hosting_networks(
+            self.ctx, ['2222'])
+        agent3 = self.plugin.get_dhcp_agents_hosting_networks(
+            self.ctx, ['3333'])
         self.assertEqual('host-c', agent1[0]['host'])
         self.assertEqual('host-c', agent2[0]['host'])
         self.assertEqual('host-d', agent3[0]['host'])