scheduler/rdb_testing_utils.py

#!/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))