• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import contextlib
6import datetime
7import unittest
8
9import common
10# import has side-effects, must appear before any django imports.
11from autotest_lib.frontend import setup_django_environment
12
13from autotest_lib.frontend import setup_test_environment
14from autotest_lib.frontend.afe import models
15
16
17class ShardHeartbeatTest(unittest.TestCase):
18    def setUp(self):
19        self._tag_creator = _TagCreator('ShardHeartbeatTest')
20        setup_test_environment.set_up()
21
22
23    def tearDown(self):
24        setup_test_environment.tear_down()
25
26
27    def testJobsWithDepsIsFilteredByShardLabel(self):
28        label = self._create_label()
29        shard = self._create_shard(label)
30        job = self._create_job_with_label(label)
31        # Should not be assigned.
32        self._create_job_with_label(self._create_label())
33
34        assigned = models.Job.assign_to_shard(shard, [])
35        self.assertEqual(set(assigned), {job})
36
37
38    def testJobsForHostsIsFilteredByShardLabel(self):
39        label = self._create_label()
40        shard = self._create_shard(label)
41        job = self._create_job_for_host(self._create_host(label))
42        # Should not be assigned.
43        self._create_job_for_host(self._create_host(self._create_label()))
44
45        assigned = models.Job.assign_to_shard(shard, [])
46        self.assertEqual(set(assigned), {job})
47
48
49    def testJobsWithKnownIDsIsIgnored(self):
50        label = self._create_label()
51        shard = self._create_shard(label)
52        known_job = self._create_job_with_label(label)
53        new_job = self._create_job_with_label(label)
54        assigned_jobs = models.Job.assign_to_shard(shard, [known_job.id])
55        self.assertEqual(set(assigned_jobs), {new_job})
56
57
58    def testOldJobsAreIgnoredWhenOptionEnabled(self):
59        with self._ignore_jobs_older_than(2):
60            label = self._create_label()
61            shard = self._create_shard(label)
62            job = self._create_job_with_label(label)
63            # Should not be assigned.
64            self._create_job_with_label(label, datetime.timedelta(hours=3))
65            assigned = models.Job.assign_to_shard(shard, [])
66            self.assertEqual(set(assigned), {job})
67
68
69    def testOldJobsAreNotIgnoredWhenOptionDisabled(self):
70        with self._ignore_jobs_older_than(0):
71            label = self._create_label()
72            shard = self._create_shard(label)
73            job = self._create_job_with_label(label,
74                                              datetime.timedelta(hours=3))
75            assigned = models.Job.assign_to_shard(shard, [])
76            self.assertEqual(set(assigned), {job})
77
78
79    @contextlib.contextmanager
80    def _ignore_jobs_older_than(self, value):
81        old = models.Job.SKIP_JOBS_CREATED_BEFORE
82        try:
83            models.Job.SKIP_JOBS_CREATED_BEFORE = value
84            yield
85        finally:
86            models.Job.SKIP_JOBS_CREATED_BEFORE = old
87
88
89    def _create_job_for_host(self, host, pending_age=None):
90        """Create a job for the given host created pending_age ago.
91
92        @param host: A models.Host object.
93        @param pending_age: A datetime.datetime object.
94        @return: A models.Job object.
95        """
96        job = models.Job.objects.create(
97            name=self._tag_creator.next(),
98            created_on=_past_timestamp(pending_age),
99        )
100        hqe = models.HostQueueEntry.objects.create(
101            job=job,
102            host_id=host.id,
103            status=models.HostQueueEntry.Status.QUEUED,
104        )
105        return job
106
107
108    def _create_job_with_label(self, label, pending_age=None):
109        """Create a job for the given label created pending_age ago.
110
111        @param host: A models.Label object.
112        @param pending_age: A datetime.datetime object.
113        @return: A models.Job object.
114        """
115        job = models.Job.objects.create(
116            name=self._tag_creator.next(),
117            created_on=_past_timestamp(pending_age),
118        )
119        job.dependency_labels.add(label)
120        hqe = models.HostQueueEntry.objects.create(
121            job=job,
122            meta_host_id=label.id,
123            status=models.HostQueueEntry.Status.QUEUED,
124        )
125        return job
126
127
128    def _create_host(self, label):
129        """Create a models.Host with the given models.Label."""
130        host = models.Host.objects.create(hostname=self._tag_creator.next())
131        host.labels.add(label)
132        return host
133
134
135    def _create_label(self):
136        """Create a models.Label."""
137        return models.Label.objects.create(name=self._tag_creator.next())
138
139
140    def _create_shard(self, label):
141        """Create a models.Shard with the givem models.Label."""
142        shard = models.Shard.objects.create(hostname=self._tag_creator.next())
143        shard.labels.add(label)
144        return shard
145
146
147class _TagCreator(object):
148    """Create unique but deterministic str tags by calling next()."""
149    def __init__(self, prefix):
150        self._prefix = prefix
151        self._count = 0
152
153    def next(self):
154        self._count += 1
155        return self._prefix + str(self._count)
156
157
158def _past_timestamp(delta=None):
159    """Compute datetime.datetime given timedelta in the past.
160
161    @param delta: A datetime.timedelta object.
162    @return: A datetime.datetime object.
163    """
164    now = datetime.datetime.now()
165    if delta is None:
166        return now
167    return now - delta
168
169
170if __name__ == '__main__':
171    unittest.main()