#!/usr/bin/python #pylint: disable-msg=C0111 # Copyright (c) 2014 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. import abc import os import common from autotest_lib.database import database_connection from autotest_lib.frontend import setup_django_environment from autotest_lib.frontend.afe import frontend_test_utils from autotest_lib.frontend.afe import models from autotest_lib.frontend.afe import rdb_model_extensions as rdb_models from autotest_lib.scheduler import monitor_db from autotest_lib.scheduler import query_managers from autotest_lib.scheduler import scheduler_lib from autotest_lib.scheduler import scheduler_models from autotest_lib.scheduler import rdb_hosts from autotest_lib.scheduler import rdb_requests from autotest_lib.server.cros import provision # Set for verbose table creation output. _DEBUG = False DEFAULT_ACLS = ['Everyone', 'my_acl'] DEFAULT_DEPS = ['a', 'b'] DEFAULT_USER = 'system' def get_default_job_params(): return {'deps': DEFAULT_DEPS, 'user': DEFAULT_USER, 'acls': DEFAULT_ACLS, 'priority': 0, 'parent_job_id': 0} def get_default_host_params(): return {'deps': DEFAULT_DEPS, 'acls': DEFAULT_ACLS} class FakeHost(rdb_hosts.RDBHost): """Fake host to use in unittests.""" def __init__(self, hostname, host_id, **kwargs): kwargs.update({'hostname': hostname, 'id': host_id}) kwargs = rdb_models.AbstractHostModel.provide_default_values( kwargs) super(FakeHost, self).__init__(**kwargs) def wire_format_response_map(response_map): wire_formatted_map = {} for request, response in response_map.iteritems(): wire_formatted_map[request] = [reply.wire_format() for reply in response] return wire_formatted_map class DBHelper(object): """Utility class for updating the database.""" def __init__(self): """Initialized django so it uses an in memory SQLite database.""" self.database = ( database_connection.TranslatingDatabase.get_test_database( translators=scheduler_lib._DB_TRANSLATORS)) self.database.connect(db_type='django') self.database.debug = _DEBUG @classmethod def get_labels(cls, **kwargs): """Get a label queryset based on the kwargs.""" return models.Label.objects.filter(**kwargs) @classmethod def get_acls(cls, **kwargs): """Get an aclgroup queryset based on the kwargs.""" return models.AclGroup.objects.filter(**kwargs) @classmethod def get_host(cls, **kwargs): """Get a host queryset based on the kwargs.""" return models.Host.objects.filter(**kwargs) @classmethod def get_hqes(cls, **kwargs): return models.HostQueueEntry.objects.filter(**kwargs) @classmethod def get_tasks(cls, **kwargs): return models.SpecialTask.objects.filter(**kwargs) @classmethod def get_shard(cls, **kwargs): return models.Shard.objects.filter(**kwargs) @classmethod def create_label(cls, name, **kwargs): label = cls.get_labels(name=name, **kwargs) return (models.Label.add_object(name=name, **kwargs) if not label else label[0]) @classmethod def create_user(cls, name): user = models.User.objects.filter(login=name) return models.User.add_object(login=name) if not user else user[0] @classmethod def create_special_task(cls, job_id=None, host_id=None, task=models.SpecialTask.Task.VERIFY, user='autotest-system'): if job_id: queue_entry = cls.get_hqes(job_id=job_id)[0] host_id = queue_entry.host.id else: queue_entry = None host = models.Host.objects.get(id=host_id) owner = cls.create_user(user) if not host: raise ValueError('Require a host to create special tasks.') return models.SpecialTask.objects.create( host=host, queue_entry=queue_entry, task=task, requested_by_id=owner.id) @classmethod def create_shard(cls, shard_hostname): """Create a shard with the given hostname if one doesn't already exist. @param shard_hostname: The hostname of the shard. """ shard = cls.get_shard(hostname=shard_hostname) return (models.Shard.objects.create(hostname=shard_hostname) if not shard else shard[0]) @classmethod def add_labels_to_host(cls, host, label_names=set([])): label_objects = set([]) for label in label_names: label_objects.add(cls.create_label(label)) host.labels.add(*label_objects) @classmethod def create_acl_group(cls, name): aclgroup = cls.get_acls(name=name) return (models.AclGroup.add_object(name=name) if not aclgroup else aclgroup[0]) @classmethod def add_deps_to_job(cls, job, dep_names=set([])): label_objects = set([]) for label in dep_names: label_objects.add(cls.create_label(label)) job.dependency_labels.add(*label_objects) @classmethod def assign_job_to_shard(cls, job_id, shard_hostname): """Assign a job to a shard. @param job: A job object without a shard. @param shard_hostname: The hostname of a shard to assign the job. @raises ValueError: If the job already has a shard. """ job_filter = models.Job.objects.filter(id=job_id, shard__isnull=True) if len(job_filter) != 1: raise ValueError('Failed to assign job %s to shard %s' % job_filter, shard_hostname) job_filter.update(shard=cls.create_shard(shard_hostname)) @classmethod def add_host_to_aclgroup(cls, host, aclgroup_names=set([])): for group_name in aclgroup_names: aclgroup = cls.create_acl_group(group_name) aclgroup.hosts.add(host) @classmethod def add_user_to_aclgroups(cls, username, aclgroup_names=set([])): user = cls.create_user(username) for group_name in aclgroup_names: aclgroup = cls.create_acl_group(group_name) aclgroup.users.add(user) @classmethod def create_host(cls, name, deps=set([]), acls=set([]), status='Ready', locked=0, lock_reason='', leased=0, protection=0, dirty=0): """Create a host. Also adds the appropriate labels to the host, and adds the host to the required acl groups. @param name: The hostname. @param kwargs: deps: The labels on the host that match job deps. acls: The aclgroups this host must be a part of. status: The status of the host. locked: 1 if the host is locked. lock_reason: non-empty string if the host is locked. leased: 1 if the host is leased. protection: Any protection level, such as Do Not Verify. dirty: 1 if the host requires cleanup. @return: The host object for the new host. """ # TODO: Modify this to use the create host request once # crbug.com/350995 is fixed. host = models.Host.add_object( hostname=name, status=status, locked=locked, lock_reason=lock_reason, leased=leased, protection=protection) cls.add_labels_to_host(host, label_names=deps) cls.add_host_to_aclgroup(host, aclgroup_names=acls) # Though we can return the host object above, this proves that the host # actually got saved in the database. For example, this will return none # if save() wasn't called on the model.Host instance. return cls.get_host(hostname=name)[0] @classmethod def update_hqe(cls, hqe_id, **kwargs): """Update the hqe with the given kwargs. @param hqe_id: The id of the hqe to update. """ models.HostQueueEntry.objects.filter(id=hqe_id).update(**kwargs) @classmethod def update_special_task(cls, task_id, **kwargs): """Update special tasks with the given kwargs. @param task_id: The if of the task to update. """ models.SpecialTask.objects.filter(id=task_id).update(**kwargs) @classmethod def add_host_to_job(cls, host, job_id, activate=0): """Add a host to the hqe of a job. @param host: An instance of the host model. @param job_id: The job to which we need to add the host. @param activate: If true, flip the active bit on the hqe. @raises ValueError: If the hqe for the job already has a host, or if the host argument isn't a Host instance. """ hqe = models.HostQueueEntry.objects.get(job_id=job_id) if hqe.host: raise ValueError('HQE for job %s already has a host' % job_id) hqe.host = host hqe.save() if activate: cls.update_hqe(hqe.id, active=True) @classmethod def increment_priority(cls, job_id): job = models.Job.objects.get(id=job_id) job.priority = job.priority + 1 job.save() class FileDatabaseHelper(object): """A helper class to setup a SQLite database backed by a file. Note that initializing a file database takes significantly longer than an in-memory database and should only be used for functional tests. """ DB_FILE = os.path.join(common.autotest_dir, 'host_scheduler_db') def initialize_database_for_testing(self, db_file_path=None): """Initialize a SQLite database for testing. To force monitor_db and the host_scheduler to use the same SQLite file database, call this method before initializing the database through frontend_test_utils. The host_scheduler is setup to look for the host_scheduler_db when invoked with --testing. @param db_file_path: The name of the file to use to create a SQLite database. Since this database is shared across different processes using a file is closer to the real world. """ if not db_file_path: db_file_path = self.DB_FILE # TODO: Move the translating database elsewhere. Monitor_db circular # imports host_scheduler. from autotest_lib.frontend import setup_test_environment from django.conf import settings self.old_django_db_name = settings.DATABASES['default']['NAME'] settings.DATABASES['default']['NAME'] = db_file_path self.db_file_path = db_file_path _db_manager = scheduler_lib.ConnectionManager(autocommit=False) _db_manager.db_connection = ( database_connection.TranslatingDatabase.get_test_database( translators=scheduler_lib._DB_TRANSLATORS)) def teardown_file_database(self): """Teardown django database settings.""" # TODO: Move the translating database elsewhere. Monitor_db circular # imports host_scheduler. from django.conf import settings settings.DATABASES['default']['NAME'] = self.old_django_db_name try: os.remove(self.db_file_path) except (OSError, AttributeError): pass class AbstractBaseRDBTester(frontend_test_utils.FrontendTestMixin): __meta__ = abc.ABCMeta _config_section = 'AUTOTEST_WEB' @staticmethod def get_request(dep_names, acl_names, priority=0, parent_job_id=0): deps = [dep.id for dep in DBHelper.get_labels(name__in=dep_names)] acls = [acl.id for acl in DBHelper.get_acls(name__in=acl_names)] return rdb_requests.AcquireHostRequest( deps=deps, acls=acls, host_id=None, priority=priority, parent_job_id=parent_job_id)._request def _release_unused_hosts(self): """Release all hosts unused by an active hqe. """ self.host_scheduler.tick() def setUp(self, inline_host_acquisition=True, setup_tables=True): """Common setup module for tests that need a jobs/host database. @param inline_host_acquisition: If True, the dispatcher tries to acquire hosts inline with the rest of the tick. """ self.db_helper = DBHelper() self._database = self.db_helper.database # Runs syncdb setting up initial database conditions self._frontend_common_setup(setup_tables=setup_tables) connection_manager = scheduler_lib.ConnectionManager(autocommit=False) self.god.stub_with(connection_manager, 'db_connection', self._database) self.god.stub_with(monitor_db, '_db_manager', connection_manager) self.god.stub_with(scheduler_models, '_db', self._database) self.god.stub_with(monitor_db, '_inline_host_acquisition', inline_host_acquisition) self._dispatcher = monitor_db.Dispatcher() self.host_scheduler = self._dispatcher._host_scheduler self.host_query_manager = query_managers.AFEHostQueryManager() self.job_query_manager = self._dispatcher._job_query_manager self._release_unused_hosts() def tearDown(self): self.god.unstub_all() self._database.disconnect() self._frontend_common_teardown() def create_job(self, user='autotest_system', deps=set([]), acls=set([]), hostless_job=False, priority=0, parent_job_id=None, shard_hostname=None): """Create a job owned by user, with the deps and acls specified. This method is a wrapper around frontend_test_utils.create_job, that also takes care of creating the appropriate deps for a job, and the appropriate acls for the given user. @raises ValueError: If no deps are specified for a job, since all jobs need at least the metahost. @raises AssertionError: If no hqe was created for the job. @return: An instance of the job model associated with the new job. """ # This is a slight hack around the implementation of # scheduler_models.is_hostless_job, even though a metahost is just # another label to the rdb. if not deps: raise ValueError('Need at least one dep for metahost') # TODO: This is a hack around the fact that frontend_test_utils still # need a metahost, but metahost is treated like any other label. metahost = self.db_helper.create_label(list(deps)[0]) job = self._create_job(metahosts=[metahost.id], priority=priority, owner=user, parent_job_id=parent_job_id) self.assert_(len(job.hostqueueentry_set.all()) == 1) self.db_helper.add_deps_to_job(job, dep_names=list(deps)[1:]) self.db_helper.add_user_to_aclgroups(user, aclgroup_names=acls) if shard_hostname: self.db_helper.assign_job_to_shard(job.id, shard_hostname) return models.Job.objects.filter(id=job.id)[0] def assert_host_db_status(self, host_id): """Assert host state right after acquisition. Call this method to check the status of any host leased by the rdb before it has been assigned to an hqe. It must be leased and ready at this point in time. @param host_id: Id of the host to check. @raises AssertionError: If the host is either not leased or Ready. """ host = models.Host.objects.get(id=host_id) self.assert_(host.leased) self.assert_(host.status == 'Ready') def check_hosts(self, host_iter): """Sanity check all hosts in the host_gen. @param host_iter: A generator/iterator of RDBClientHostWrappers. eg: The generator returned by rdb_lib.acquire_hosts. If a request was not satisfied this iterator can contain None. @raises AssertionError: If any of the sanity checks fail. """ for host in host_iter: if host: self.assert_host_db_status(host.id) self.assert_(host.leased == 1) def create_suite(self, user='autotest_system', num=2, priority=0, board='z', build='x', acls=set()): """Create num jobs with the same parent_job_id, board, build, priority. @return: A dictionary with the parent job object keyed as 'parent_job' and all other jobs keyed at an index from 0-num. """ jobs = {} # Create a hostless parent job without an hqe or deps. Since the # hostless job does nothing, we need to hand craft cros-version. parent_job = self._create_job(owner=user, priority=priority) jobs['parent_job'] = parent_job build = '%s:%s' % (provision.CROS_VERSION_PREFIX, build) for job_index in range(0, num): jobs[job_index] = self.create_job(user=user, priority=priority, deps=set([board, build]), acls=acls, parent_job_id=parent_job.id) return jobs def check_host_assignment(self, job_id, host_id): """Check is a job<->host assignment is valid. Uses the deps of a job and the aclgroups the owner of the job is in to see if the given host can be used to run the given job. Also checks that the host-job assignment has Not been made, but that the host is no longer in the available hosts pool. Use this method to check host assignements made by the rdb, Before they're handed off to the scheduler, since the scheduler. @param job_id: The id of the job to use in the compatibility check. @param host_id: The id of the host to check for compatibility. @raises AssertionError: If the job and the host are incompatible. """ job = models.Job.objects.get(id=job_id) host = models.Host.objects.get(id=host_id) hqe = job.hostqueueentry_set.all()[0] # Confirm that the host has not been assigned, either to another hqe # or the this one. all_hqes = models.HostQueueEntry.objects.filter( host_id=host_id, complete=0) self.assert_(len(all_hqes) <= 1) self.assert_(hqe.host_id == None) self.assert_host_db_status(host_id) # Assert that all deps of the job are satisfied. job_deps = set([d.name for d in job.dependency_labels.all()]) host_labels = set([l.name for l in host.labels.all()]) self.assert_(job_deps.intersection(host_labels) == job_deps) # Assert that the owner of the job is in at least one of the # groups that owns the host. job_owner_aclgroups = set([job_acl.name for job_acl in job.user().aclgroup_set.all()]) host_aclgroups = set([host_acl.name for host_acl in host.aclgroup_set.all()]) self.assert_(job_owner_aclgroups.intersection(host_aclgroups))