1#!/usr/bin/env python
2#
3# Copyright (C) 2018 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18import datetime
19import unittest
20
21try:
22    from unittest import mock
23except ImportError:
24    import mock
25
26from webapp.src import vtslab_status as Status
27from webapp.src.proto import model
28from webapp.src.scheduler import schedule_worker
29from webapp.src.testing import unittest_base
30from webapp.src.utils import model_util
31
32
33class ModelTest(unittest_base.UnitTestBase):
34    """Tests for PeriodicJobHeartBeat cron class."""
35
36    def testJobAndScheduleModel(self):
37        """Asserts JobModel and ScheduleModel.
38
39        When JobModel's status is changed, ScheduleModel's error_count is
40        changed based on the status. This should not be applied before JobModel
41        entity is updated to Datastore.
42        """
43        period = 360
44
45        lab = self.GenerateLabModel()
46        lab.put()
47
48        device = self.GenerateDeviceModel(hostname=lab.hostname)
49        device.put()
50
51        schedule = self.GenerateScheduleModel(
52            device_model=device, lab_model=lab, period=period)
53        schedule.put()
54
55        build_dict = self.GenerateBuildModel(schedule)
56        for key in build_dict:
57            build_dict[key].put()
58
59        # Mocking ScheduleHandler and essential methods.
60        scheduler = schedule_worker.ScheduleHandler(mock.Mock())
61        scheduler.response = mock.Mock()
62        scheduler.response.write = mock.Mock()
63        scheduler.request.get = mock.MagicMock(return_value="")
64
65        print("\nCreating a job...")
66        scheduler.post()
67        jobs = model.JobModel.query().fetch()
68        self.assertEqual(1, len(jobs))
69
70        print("Occurring infra error...")
71        job = jobs[0]
72        job.status = Status.JOB_STATUS_DICT["infra-err"]
73        parent_schedule = job.parent_schedule.get()
74        parent_from_db = model.ScheduleModel.query().fetch()[0]
75
76        # in test error_count could be None but in real there will be no None.
77        self.assertNotEqual(1, parent_schedule.error_count)
78        self.assertNotEqual(1, parent_from_db.error_count)
79
80        # error count should be changed after put
81        job.put()
82        model_util.UpdateParentSchedule(job, job.status)
83        self.assertEqual(1, parent_schedule.error_count)
84        self.assertEqual(1, parent_from_db.error_count)
85
86        print("Suspending a job...")
87        for num in xrange(2):
88            jobs = model.JobModel.query().fetch()
89            for job in jobs:
90                job.timestamp = datetime.datetime.now() - datetime.timedelta(
91                    minutes=(period + 10))
92                job.put()
93
94            parent_from_db = model.ScheduleModel.query().fetch()[0]
95            self.assertEqual(1 + num, parent_schedule.error_count)
96            self.assertEqual(1 + num, parent_from_db.error_count)
97
98            # reset a device manually to re-schedule
99            device = model.DeviceModel.query().fetch()[0]
100            device.status = Status.DEVICE_STATUS_DICT["fastboot"]
101            device.scheduling_status = (
102                Status.DEVICE_SCHEDULING_STATUS_DICT["free"])
103            device.timestamp = datetime.datetime.now()
104            device.put()
105
106            scheduler.post()
107            jobs = model.JobModel.query().fetch()
108            self.assertEqual(2 + num, len(jobs))
109
110            ready_jobs = model.JobModel.query(
111                model.JobModel.status == Status.JOB_STATUS_DICT[
112                    "ready"]).fetch()
113            self.assertEqual(1, len(ready_jobs))
114
115            ready_job = ready_jobs[0]
116            ready_job.status = Status.JOB_STATUS_DICT["infra-err"]
117            parent_schedule = ready_job.parent_schedule.get()
118            parent_from_db = model.ScheduleModel.query().fetch()[0]
119            self.assertEqual(1 + num, parent_schedule.error_count)
120            self.assertEqual(1 + num, parent_from_db.error_count)
121
122            # # error count should be changed after put
123            ready_job.put()
124            model_util.UpdateParentSchedule(ready_job, ready_job.status)
125            self.assertEqual(2 + num, parent_schedule.error_count)
126            self.assertEqual(2 + num, parent_from_db.error_count)
127
128        print("Asserting a schedule's suspend status...")
129        # after three errors the schedule should be suspended.
130        schedule_from_db = model.ScheduleModel.query().fetch()[0]
131        schedule_from_db.put()
132        self.assertEqual(3, schedule_from_db.error_count)
133        self.assertEqual(True, schedule_from_db.suspended)
134
135        # reset a device manually to re-schedule
136        device = model.DeviceModel.query().fetch()[0]
137        device.status = Status.DEVICE_STATUS_DICT["fastboot"]
138        device.scheduling_status = (
139            Status.DEVICE_SCHEDULING_STATUS_DICT["free"])
140        device.timestamp = datetime.datetime.now()
141        device.put()
142
143        print("Asserting that job creation is blocked...")
144        jobs = model.JobModel.query().fetch()
145        self.assertEqual(3, len(jobs))
146
147        for job in jobs:
148            job.timestamp = datetime.datetime.now() - datetime.timedelta(
149                minutes=(period + 10))
150            job.put()
151
152        scheduler.post()
153
154        # a job should not be created.
155        jobs = model.JobModel.query().fetch()
156        self.assertEqual(3, len(jobs))
157
158        print("Asserting that job creation is allowed after resuming...")
159        schedule_from_db = model.ScheduleModel.query().fetch()[0]
160        schedule_from_db.suspended = False
161        schedule_from_db.put()
162
163        scheduler.post()
164
165        jobs = model.JobModel.query().fetch()
166        self.assertEqual(4, len(jobs))
167
168
169if __name__ == "__main__":
170    unittest.main()
171