#!/usr/bin/python # # Copyright (c) 2012 The Chromium OS Authors. All rights reserved. # Use of this source code is governed by a BSD-style license that can be # found in the LICENSE file. """Unit tests for frontend/afe/site_rpc_interface.py.""" import __builtin__ import datetime import mox import StringIO import unittest import common from autotest_lib.frontend import setup_django_environment from autotest_lib.frontend.afe import frontend_test_utils from autotest_lib.frontend.afe import models, model_logic, rpc_utils from autotest_lib.client.common_lib import control_data, error from autotest_lib.client.common_lib import global_config from autotest_lib.client.common_lib import lsbrelease_utils from autotest_lib.client.common_lib import priorities from autotest_lib.client.common_lib.cros import dev_server from autotest_lib.frontend.afe import rpc_interface, site_rpc_interface from autotest_lib.server import utils from autotest_lib.server.cros.dynamic_suite import control_file_getter from autotest_lib.server.cros.dynamic_suite import constants from autotest_lib.server.hosts import moblab_host CLIENT = control_data.CONTROL_TYPE_NAMES.CLIENT SERVER = control_data.CONTROL_TYPE_NAMES.SERVER class SiteRpcInterfaceTest(mox.MoxTestBase, frontend_test_utils.FrontendTestMixin): """Unit tests for functions in site_rpc_interface.py. @var _NAME: fake suite name. @var _BOARD: fake board to reimage. @var _BUILD: fake build with which to reimage. @var _PRIORITY: fake priority with which to reimage. """ _NAME = 'name' _BOARD = 'link' _BUILD = 'link-release/R36-5812.0.0' _PRIORITY = priorities.Priority.DEFAULT _TIMEOUT = 24 def setUp(self): super(SiteRpcInterfaceTest, self).setUp() self._SUITE_NAME = site_rpc_interface.canonicalize_suite_name( self._NAME) self.dev_server = self.mox.CreateMock(dev_server.ImageServer) self._frontend_common_setup(fill_data=False) def tearDown(self): self._frontend_common_teardown() def _setupDevserver(self): self.mox.StubOutClassWithMocks(dev_server, 'ImageServer') dev_server.ImageServer.resolve(self._BUILD).AndReturn(self.dev_server) def _mockDevServerGetter(self, get_control_file=True): self._setupDevserver() if get_control_file: self.getter = self.mox.CreateMock( control_file_getter.DevServerGetter) self.mox.StubOutWithMock(control_file_getter.DevServerGetter, 'create') control_file_getter.DevServerGetter.create( mox.IgnoreArg(), mox.IgnoreArg()).AndReturn(self.getter) def _mockRpcUtils(self, to_return, control_file_substring=''): """Fake out the autotest rpc_utils module with a mockable class. @param to_return: the value that rpc_utils.create_job_common() should be mocked out to return. @param control_file_substring: A substring that is expected to appear in the control file output string that is passed to create_job_common. Default: '' """ download_started_time = constants.DOWNLOAD_STARTED_TIME payload_finished_time = constants.PAYLOAD_FINISHED_TIME self.mox.StubOutWithMock(rpc_utils, 'create_job_common') rpc_utils.create_job_common(mox.And(mox.StrContains(self._NAME), mox.StrContains(self._BUILD)), priority=self._PRIORITY, timeout_mins=self._TIMEOUT*60, max_runtime_mins=self._TIMEOUT*60, control_type='Server', control_file=mox.And(mox.StrContains(self._BOARD), mox.StrContains(self._BUILD), mox.StrContains( control_file_substring)), hostless=True, keyvals=mox.And(mox.In(download_started_time), mox.In(payload_finished_time)) ).AndReturn(to_return) def testStageBuildFail(self): """Ensure that a failure to stage the desired build fails the RPC.""" self._setupDevserver() self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.dev_server.stage_artifacts( self._BUILD, ['test_suites']).AndRaise( dev_server.DevServerException()) self.mox.ReplayAll() self.assertRaises(error.StageControlFileFailure, site_rpc_interface.create_suite_job, name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None) def testGetControlFileFail(self): """Ensure that a failure to get needed control file fails the RPC.""" self._mockDevServerGetter() self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.dev_server.stage_artifacts(self._BUILD, ['test_suites']).AndReturn(True) self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.getter.get_control_file_contents_by_name( self._SUITE_NAME).AndReturn(None) self.mox.ReplayAll() self.assertRaises(error.ControlFileEmpty, site_rpc_interface.create_suite_job, name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None) def testGetControlFileListFail(self): """Ensure that a failure to get needed control file fails the RPC.""" self._mockDevServerGetter() self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.dev_server.stage_artifacts(self._BUILD, ['test_suites']).AndReturn(True) self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.getter.get_control_file_contents_by_name( self._SUITE_NAME).AndRaise(error.NoControlFileList()) self.mox.ReplayAll() self.assertRaises(error.NoControlFileList, site_rpc_interface.create_suite_job, name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None) def testBadNumArgument(self): """Ensure we handle bad values for the |num| argument.""" self.assertRaises(error.SuiteArgumentException, site_rpc_interface.create_suite_job, name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None, num='goo') self.assertRaises(error.SuiteArgumentException, site_rpc_interface.create_suite_job, name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None, num=[]) self.assertRaises(error.SuiteArgumentException, site_rpc_interface.create_suite_job, name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None, num='5') def testCreateSuiteJobFail(self): """Ensure that failure to schedule the suite job fails the RPC.""" self._mockDevServerGetter() self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.dev_server.stage_artifacts(self._BUILD, ['test_suites']).AndReturn(True) self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.getter.get_control_file_contents_by_name( self._SUITE_NAME).AndReturn('f') self.dev_server.url().AndReturn('mox_url') self._mockRpcUtils(-1) self.mox.ReplayAll() self.assertEquals( site_rpc_interface.create_suite_job(name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None), -1) def testCreateSuiteJobSuccess(self): """Ensures that success results in a successful RPC.""" self._mockDevServerGetter() self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.dev_server.stage_artifacts(self._BUILD, ['test_suites']).AndReturn(True) self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.getter.get_control_file_contents_by_name( self._SUITE_NAME).AndReturn('f') self.dev_server.url().AndReturn('mox_url') job_id = 5 self._mockRpcUtils(job_id) self.mox.ReplayAll() self.assertEquals( site_rpc_interface.create_suite_job(name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None), job_id) def testCreateSuiteJobNoHostCheckSuccess(self): """Ensures that success results in a successful RPC.""" self._mockDevServerGetter() self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.dev_server.stage_artifacts(self._BUILD, ['test_suites']).AndReturn(True) self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.getter.get_control_file_contents_by_name( self._SUITE_NAME).AndReturn('f') self.dev_server.url().AndReturn('mox_url') job_id = 5 self._mockRpcUtils(job_id) self.mox.ReplayAll() self.assertEquals( site_rpc_interface.create_suite_job(name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None, check_hosts=False), job_id) def testCreateSuiteIntegerNum(self): """Ensures that success results in a successful RPC.""" self._mockDevServerGetter() self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.dev_server.stage_artifacts(self._BUILD, ['test_suites']).AndReturn(True) self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.getter.get_control_file_contents_by_name( self._SUITE_NAME).AndReturn('f') self.dev_server.url().AndReturn('mox_url') job_id = 5 self._mockRpcUtils(job_id, control_file_substring='num=17') self.mox.ReplayAll() self.assertEquals( site_rpc_interface.create_suite_job(name=self._NAME, board=self._BOARD, build=self._BUILD, pool=None, check_hosts=False, num=17), job_id) def testCreateSuiteJobControlFileSupplied(self): """Ensure we can supply the control file to create_suite_job.""" self._mockDevServerGetter(get_control_file=False) self.dev_server.url().AndReturn('mox_url') self.dev_server.get_server_name(mox.IgnoreArg()).AndReturn('mox_url') self.dev_server.stage_artifacts(self._BUILD, ['test_suites']).AndReturn(True) self.dev_server.url().AndReturn('mox_url') job_id = 5 self._mockRpcUtils(job_id) self.mox.ReplayAll() self.assertEquals( site_rpc_interface.create_suite_job(name='%s/%s' % (self._NAME, self._BUILD), board=None, build=self._BUILD, pool=None, control_file='CONTROL FILE'), job_id) def setIsMoblab(self, is_moblab): """Set utils.is_moblab result. @param is_moblab: Value to have utils.is_moblab to return. """ self.mox.StubOutWithMock(utils, 'is_moblab') utils.is_moblab().AndReturn(is_moblab) def testMoblabOnlyDecorator(self): """Ensure the moblab only decorator gates functions properly.""" self.setIsMoblab(False) self.mox.ReplayAll() self.assertRaises(error.RPCException, site_rpc_interface.get_config_values) def testGetConfigValues(self): """Ensure that the config object is properly converted to a dict.""" self.setIsMoblab(True) config_mock = self.mox.CreateMockAnything() site_rpc_interface._CONFIG = config_mock config_mock.get_sections().AndReturn(['section1', 'section2']) config_mock.config = self.mox.CreateMockAnything() config_mock.config.items('section1').AndReturn([('item1', 'value1'), ('item2', 'value2')]) config_mock.config.items('section2').AndReturn([('item3', 'value3'), ('item4', 'value4')]) rpc_utils.prepare_for_serialization( {'section1' : [('item1', 'value1'), ('item2', 'value2')], 'section2' : [('item3', 'value3'), ('item4', 'value4')]}) self.mox.ReplayAll() site_rpc_interface.get_config_values() def _mockReadFile(self, path, lines=[]): """Mock out reading a file line by line. @param path: Path of the file we are mock reading. @param lines: lines of the mock file that will be returned when readLine() is called. """ mockFile = self.mox.CreateMockAnything() for line in lines: mockFile.readline().AndReturn(line) mockFile.readline() mockFile.close() open(path).AndReturn(mockFile) def testUpdateConfig(self): """Ensure that updating the config works as expected.""" self.setIsMoblab(True) site_rpc_interface.os = self.mox.CreateMockAnything() self.mox.StubOutWithMock(__builtin__, 'open') self._mockReadFile(global_config.DEFAULT_CONFIG_FILE) self.mox.StubOutWithMock(lsbrelease_utils, 'is_moblab') lsbrelease_utils.is_moblab().AndReturn(True) self._mockReadFile(global_config.DEFAULT_MOBLAB_FILE, ['[section1]', 'item1: value1']) site_rpc_interface.os = self.mox.CreateMockAnything() site_rpc_interface.os.path = self.mox.CreateMockAnything() site_rpc_interface.os.path.exists( site_rpc_interface._CONFIG.shadow_file).AndReturn( True) mockShadowFile = self.mox.CreateMockAnything() mockShadowFileContents = StringIO.StringIO() mockShadowFile.__enter__().AndReturn(mockShadowFileContents) mockShadowFile.__exit__(mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()) open(site_rpc_interface._CONFIG.shadow_file, 'w').AndReturn(mockShadowFile) site_rpc_interface.os.system('sudo reboot') self.mox.ReplayAll() site_rpc_interface.update_config_handler( {'section1' : [('item1', 'value1'), ('item2', 'value2')], 'section2' : [('item3', 'value3'), ('item4', 'value4')]}) # item1 should not be in the new shadow config as its updated value # matches the original config's value. self.assertEquals( mockShadowFileContents.getvalue(), '[section2]\nitem3 = value3\nitem4 = value4\n\n' '[section1]\nitem2 = value2\n\n') def testResetConfig(self): """Ensure that reset opens the shadow_config file for writing.""" self.setIsMoblab(True) config_mock = self.mox.CreateMockAnything() site_rpc_interface._CONFIG = config_mock config_mock.shadow_file = 'shadow_config.ini' self.mox.StubOutWithMock(__builtin__, 'open') mockFile = self.mox.CreateMockAnything() file_contents = self.mox.CreateMockAnything() mockFile.__enter__().AndReturn(file_contents) mockFile.__exit__(mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()) open(config_mock.shadow_file, 'w').AndReturn(mockFile) site_rpc_interface.os = self.mox.CreateMockAnything() site_rpc_interface.os.system('sudo reboot') self.mox.ReplayAll() site_rpc_interface.reset_config_settings() def testSetBotoKey(self): """Ensure that the botokey path supplied is copied correctly.""" self.setIsMoblab(True) boto_key = '/tmp/boto' site_rpc_interface.os.path = self.mox.CreateMockAnything() site_rpc_interface.os.path.exists(boto_key).AndReturn( True) site_rpc_interface.shutil = self.mox.CreateMockAnything() site_rpc_interface.shutil.copyfile( boto_key, site_rpc_interface.MOBLAB_BOTO_LOCATION) self.mox.ReplayAll() site_rpc_interface.set_boto_key(boto_key) def testSetLaunchControlKey(self): """Ensure that the Launch Control key path supplied is copied correctly. """ self.setIsMoblab(True) launch_control_key = '/tmp/launch_control' site_rpc_interface.os = self.mox.CreateMockAnything() site_rpc_interface.os.path = self.mox.CreateMockAnything() site_rpc_interface.os.path.exists(launch_control_key).AndReturn( True) site_rpc_interface.shutil = self.mox.CreateMockAnything() site_rpc_interface.shutil.copyfile( launch_control_key, moblab_host.MOBLAB_LAUNCH_CONTROL_KEY_LOCATION) site_rpc_interface.os.system('sudo restart moblab-devserver-init') self.mox.ReplayAll() site_rpc_interface.set_launch_control_key(launch_control_key) def _get_records_for_sending_to_master(self): return [{'control_file': 'foo', 'control_type': 1, 'created_on': datetime.datetime(2014, 8, 21), 'drone_set': None, 'email_list': '', 'max_runtime_hrs': 72, 'max_runtime_mins': 1440, 'name': 'dummy', 'owner': 'autotest_system', 'parse_failed_repair': True, 'priority': 40, 'reboot_after': 0, 'reboot_before': 1, 'run_reset': True, 'run_verify': False, 'synch_count': 0, 'test_retry': 10, 'timeout': 24, 'timeout_mins': 1440, 'id': 1 }], [{ 'aborted': False, 'active': False, 'complete': False, 'deleted': False, 'execution_subdir': '', 'finished_on': None, 'started_on': None, 'status': 'Queued', 'id': 1 }] def _do_heartbeat_and_assert_response(self, shard_hostname='shard1', upload_jobs=(), upload_hqes=(), known_jobs=(), known_hosts=(), **kwargs): known_job_ids = [job.id for job in known_jobs] known_host_ids = [host.id for host in known_hosts] known_host_statuses = [host.status for host in known_hosts] retval = site_rpc_interface.shard_heartbeat( shard_hostname=shard_hostname, jobs=upload_jobs, hqes=upload_hqes, known_job_ids=known_job_ids, known_host_ids=known_host_ids, known_host_statuses=known_host_statuses) self._assert_shard_heartbeat_response(shard_hostname, retval, **kwargs) return shard_hostname def _assert_shard_heartbeat_response(self, shard_hostname, retval, jobs=[], hosts=[], hqes=[]): retval_hosts, retval_jobs = retval['hosts'], retval['jobs'] expected_jobs = [ (job.id, job.name, shard_hostname) for job in jobs] returned_jobs = [(job['id'], job['name'], job['shard']['hostname']) for job in retval_jobs] self.assertEqual(returned_jobs, expected_jobs) expected_hosts = [(host.id, host.hostname) for host in hosts] returned_hosts = [(host['id'], host['hostname']) for host in retval_hosts] self.assertEqual(returned_hosts, expected_hosts) retval_hqes = [] for job in retval_jobs: retval_hqes += job['hostqueueentry_set'] expected_hqes = [(hqe.id) for hqe in hqes] returned_hqes = [(hqe['id']) for hqe in retval_hqes] self.assertEqual(returned_hqes, expected_hqes) def _send_records_to_master_helper( self, jobs, hqes, shard_hostname='host1', exception_to_throw=error.UnallowedRecordsSentToMaster, aborted=False): job_id = rpc_interface.create_job(name='dummy', priority='Medium', control_file='foo', control_type=SERVER, test_retry=10, hostless=True) job = models.Job.objects.get(pk=job_id) shard = models.Shard.objects.create(hostname='host1') job.shard = shard job.save() if aborted: job.hostqueueentry_set.update(aborted=True) job.shard = None job.save() hqe = job.hostqueueentry_set.all()[0] if not exception_to_throw: self._do_heartbeat_and_assert_response( shard_hostname=shard_hostname, upload_jobs=jobs, upload_hqes=hqes) else: self.assertRaises( exception_to_throw, self._do_heartbeat_and_assert_response, shard_hostname=shard_hostname, upload_jobs=jobs, upload_hqes=hqes) def testSendingRecordsToMaster(self): """Send records to the master and ensure they are persisted.""" jobs, hqes = self._get_records_for_sending_to_master() hqes[0]['status'] = 'Completed' self._send_records_to_master_helper( jobs=jobs, hqes=hqes, exception_to_throw=None) # Check the entry was actually written to db self.assertEqual(models.HostQueueEntry.objects.all()[0].status, 'Completed') def testSendingRecordsToMasterAbortedOnMaster(self): """Send records to the master and ensure they are persisted.""" jobs, hqes = self._get_records_for_sending_to_master() hqes[0]['status'] = 'Completed' self._send_records_to_master_helper( jobs=jobs, hqes=hqes, exception_to_throw=None, aborted=True) # Check the entry was actually written to db self.assertEqual(models.HostQueueEntry.objects.all()[0].status, 'Completed') def testSendingRecordsToMasterJobAssignedToDifferentShard(self): """Ensure records that belong to a different shard are rejected.""" jobs, hqes = self._get_records_for_sending_to_master() models.Shard.objects.create(hostname='other_shard') self._send_records_to_master_helper( jobs=jobs, hqes=hqes, shard_hostname='other_shard') def testSendingRecordsToMasterJobHqeWithoutJob(self): """Ensure update for hqe without update for it's job gets rejected.""" _, hqes = self._get_records_for_sending_to_master() self._send_records_to_master_helper( jobs=[], hqes=hqes) def testSendingRecordsToMasterNotExistingJob(self): """Ensure update for non existing job gets rejected.""" jobs, hqes = self._get_records_for_sending_to_master() jobs[0]['id'] = 3 self._send_records_to_master_helper( jobs=jobs, hqes=hqes) def _createShardAndHostWithLabel(self, shard_hostname='shard1', host_hostname='host1', label_name='board:lumpy'): label = models.Label.objects.create(name=label_name) shard = models.Shard.objects.create(hostname=shard_hostname) shard.labels.add(label) host = models.Host.objects.create(hostname=host_hostname, leased=False) host.labels.add(label) return shard, host, label def _createJobForLabel(self, label): job_id = rpc_interface.create_job(name='dummy', priority='Medium', control_file='foo', control_type=CLIENT, meta_hosts=[label.name], dependencies=(label.name,)) return models.Job.objects.get(id=job_id) def testShardHeartbeatFetchHostlessJob(self): """Create a hostless job and ensure it's not assigned to a shard.""" shard1, host1, lumpy_label = self._createShardAndHostWithLabel( 'shard1', 'host1', 'board:lumpy') label2 = models.Label.objects.create(name='bluetooth', platform=False) job1 = self._create_job(hostless=True) # Hostless jobs should be executed by the global scheduler. self._do_heartbeat_and_assert_response(hosts=[host1]) def testShardRetrieveJobs(self): """Create jobs and retrieve them.""" # should never be returned by heartbeat leased_host = models.Host.objects.create(hostname='leased_host', leased=True) shard1, host1, lumpy_label = self._createShardAndHostWithLabel() shard2, host2, grumpy_label = self._createShardAndHostWithLabel( 'shard2', 'host2', 'board:grumpy') leased_host.labels.add(lumpy_label) job1 = self._createJobForLabel(lumpy_label) job2 = self._createJobForLabel(grumpy_label) job_completed = self._createJobForLabel(lumpy_label) # Job is already being run, so don't sync it job_completed.hostqueueentry_set.update(complete=True) job_completed.hostqueueentry_set.create(complete=False) job_active = self._createJobForLabel(lumpy_label) # Job is already started, so don't sync it job_active.hostqueueentry_set.update(active=True) job_active.hostqueueentry_set.create(complete=False, active=False) self._do_heartbeat_and_assert_response( jobs=[job1], hosts=[host1], hqes=job1.hostqueueentry_set.all()) self._do_heartbeat_and_assert_response( shard_hostname=shard2.hostname, jobs=[job2], hosts=[host2], hqes=job2.hostqueueentry_set.all()) host3 = models.Host.objects.create(hostname='host3', leased=False) host3.labels.add(lumpy_label) self._do_heartbeat_and_assert_response( known_jobs=[job1], known_hosts=[host1], hosts=[host3]) def testResendJobsAfterFailedHeartbeat(self): """Create jobs, retrieve them, fail on client, fetch them again.""" shard1, host1, lumpy_label = self._createShardAndHostWithLabel() job1 = self._createJobForLabel(lumpy_label) self._do_heartbeat_and_assert_response( jobs=[job1], hqes=job1.hostqueueentry_set.all(), hosts=[host1]) # Make sure it's resubmitted by sending last_job=None again self._do_heartbeat_and_assert_response( known_hosts=[host1], jobs=[job1], hqes=job1.hostqueueentry_set.all(), hosts=[]) # Now it worked, make sure it's not sent again self._do_heartbeat_and_assert_response( known_jobs=[job1], known_hosts=[host1]) job1 = models.Job.objects.get(pk=job1.id) job1.hostqueueentry_set.all().update(complete=True) # Job is completed, make sure it's not sent again self._do_heartbeat_and_assert_response( known_hosts=[host1]) job2 = self._createJobForLabel(lumpy_label) # job2's creation was later, it should be returned now. self._do_heartbeat_and_assert_response( known_hosts=[host1], jobs=[job2], hqes=job2.hostqueueentry_set.all()) self._do_heartbeat_and_assert_response( known_jobs=[job2], known_hosts=[host1]) job2 = models.Job.objects.get(pk=job2.pk) job2.hostqueueentry_set.update(aborted=True) # Setting a job to a complete status will set the shard_id to None in # scheduler_models. We have to emulate that here, because we use Django # models in tests. job2.shard = None job2.save() self._do_heartbeat_and_assert_response( known_jobs=[job2], known_hosts=[host1], jobs=[job2], hqes=job2.hostqueueentry_set.all()) site_rpc_interface.delete_shard(hostname=shard1.hostname) self.assertRaises( models.Shard.DoesNotExist, models.Shard.objects.get, pk=shard1.id) job1 = models.Job.objects.get(pk=job1.id) lumpy_label = models.Label.objects.get(pk=lumpy_label.id) host1 = models.Host.objects.get(pk=host1.id) self.assertIsNone(job1.shard) self.assertEqual(len(lumpy_label.shard_set.all()), 0) self.assertIsNone(host1.shard) self.assertEqual([s.task for s in host1.specialtask_set.all()], ['Repair']) def testCreateListShard(self): """Retrieve a list of all shards.""" lumpy_label = models.Label.objects.create(name='board:lumpy', platform=True) stumpy_label = models.Label.objects.create(name='board:stumpy', platform=True) shard_id = site_rpc_interface.add_shard( hostname='host1', labels='board:lumpy,board:stumpy') self.assertRaises(model_logic.ValidationError, site_rpc_interface.add_shard, hostname='host1', labels='board:lumpy,board:stumpy') shard = models.Shard.objects.get(pk=shard_id) self.assertEqual(shard.hostname, 'host1') self.assertEqual(shard.labels.values_list('pk')[0], (lumpy_label.id,)) self.assertEqual(shard.labels.values_list('pk')[1], (stumpy_label.id,)) self.assertEqual(site_rpc_interface.get_shards(), [{'labels': ['board:lumpy','board:stumpy'], 'hostname': 'host1', 'id': 1}]) def testResendHostsAfterFailedHeartbeat(self): """Check that master accepts resending updated records after failure.""" shard1, host1, lumpy_label = self._createShardAndHostWithLabel() # Send the host self._do_heartbeat_and_assert_response(hosts=[host1]) # Send it again because previous one didn't persist correctly self._do_heartbeat_and_assert_response(hosts=[host1]) # Now it worked, make sure it isn't sent again self._do_heartbeat_and_assert_response(known_hosts=[host1]) if __name__ == '__main__': unittest.main()