1# 2# Copyright (C) 2018 The Android Open Source Project 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15# 16 17import datetime 18import random 19import string 20import unittest 21 22try: 23 from unittest import mock 24except ImportError: 25 import mock 26 27from google.appengine.ext import ndb 28from google.appengine.ext import testbed 29 30from webapp.src import vtslab_status as Status 31from webapp.src.proto import model 32 33 34class UnitTestBase(unittest.TestCase): 35 """Base class for unittest. 36 37 Attributes: 38 testbed: A Testbed instance which provides local unit testing. 39 random_strs: a list of strings generated by GetRandomString() method 40 in order to avoid duplicates. 41 """ 42 random_strs = [] 43 44 def setUp(self): 45 """Initializes unittest.""" 46 # Create the Testbed class instance and initialize service stubs. 47 self.testbed = testbed.Testbed() 48 self.testbed.activate() 49 self.testbed.init_datastore_v3_stub() 50 self.testbed.init_memcache_stub() 51 self.testbed.init_mail_stub() 52 self.testbed.setup_env(app_id="vtslab-schedule-unittest") 53 # Clear cache between tests. 54 ndb.get_context().clear_cache() 55 56 def tearDown(self): 57 self.testbed.deactivate() 58 59 def GetRandomString(self, length=7): 60 """Generates and returns a random string. 61 62 Args: 63 length: an integer, string length. 64 65 Returns: 66 a random string. 67 """ 68 new_str = "" 69 while new_str == "" or new_str in self.random_strs: 70 new_str = "".join( 71 random.choice(string.ascii_letters + string.digits) 72 for _ in range(length)) 73 return new_str 74 75 def GenerateLabModel(self, lab_name=None, host_name=None): 76 """Builds model.LabModel with given information. 77 78 Args: 79 lab_name: a string, lab name. 80 host_name: a string, host name. 81 82 Returns: 83 model.LabModel instance. 84 """ 85 lab = model.LabModel() 86 lab.name = lab_name if lab_name else self.GetRandomString() 87 lab.hostname = host_name if host_name else self.GetRandomString() 88 lab.owner = "test@abc.com" 89 lab.ip = "100.100.100.100" 90 return lab 91 92 def GenerateDeviceModel( 93 self, 94 status=Status.DEVICE_STATUS_DICT["fastboot"], 95 scheduling_status=Status.DEVICE_SCHEDULING_STATUS_DICT["free"], 96 **kwargs): 97 """Builds model.DeviceModel with given information. 98 99 Args: 100 status: an integer, device's initial status. 101 scheduling_status: an integer, device's initial scheduling status. 102 **kwargs: the optional arguments. 103 104 Returns: 105 model.DeviceModel instance. 106 """ 107 device = model.DeviceModel() 108 device.status = status 109 device.scheduling_status = scheduling_status 110 device.timestamp = datetime.datetime.now() 111 112 skip_list = ["status", "scheduling_status", "timestamp"] 113 set_or_empty = [] 114 for arg in device._properties: 115 if arg in skip_list or (arg in set_or_empty and arg not in kwargs): 116 continue 117 if arg in kwargs: 118 value = kwargs[arg] 119 elif isinstance(device._properties[arg], ndb.StringProperty): 120 value = self.GetRandomString() 121 elif isinstance(device._properties[arg], ndb.IntegerProperty): 122 value = 0 123 elif isinstance(device._properties[arg], ndb.BooleanProperty): 124 value = False 125 else: 126 print("A type of property '{}' is not supported.".format(arg)) 127 continue 128 if device._properties[arg]._repeated and type(value) is not list: 129 value = [value] 130 setattr(device, arg, value) 131 return device 132 133 def GenerateScheduleModel( 134 self, 135 device_model=None, 136 lab_model=None, 137 priority="medium", 138 period=360, 139 retry_count=1, 140 shards=1, 141 lab_name=None, 142 device_storage_type=Status.STORAGE_TYPE_DICT["PAB"], 143 device_branch=None, 144 device_target=None, 145 gsi_storage_type=Status.STORAGE_TYPE_DICT["PAB"], 146 gsi_build_target=None, 147 test_storage_type=Status.STORAGE_TYPE_DICT["PAB"], 148 test_build_target=None, 149 required_signed_device_build=False, 150 **kwargs): 151 """Builds model.ScheduleModel with given information. 152 153 Args: 154 device_model: a model.DeviceModel instance to refer device product. 155 lab_model: a model.LabModel instance to refer host name. 156 priority: a string, scheduling priority 157 period: an integer, scheduling period. 158 retry_count: an integer, scheduling retry count. 159 shards: an integer, # ways of device shards. 160 lab_name: a string, target lab name. 161 device_storage_type: an integer, device storage type 162 device_branch: a string, device build branch. 163 device_target: a string, device build target. 164 gsi_storage_type: an integer, GSI storage type 165 gsi_build_target: a string, GSI build target. 166 test_storage_type: an integer, test storage type 167 test_build_target: a string, test build target. 168 required_signed_device_build: a boolean, True to schedule for signed 169 device build, False if not. 170 **kwargs: the optional arguments. 171 172 Returns: 173 model.ScheduleModel instance. 174 """ 175 176 if device_model: 177 device_product = device_model.product 178 device_target = self.GetRandomString(4) 179 elif device_target: 180 device_product, device_target = device_target.split("-") 181 else: 182 device_product = self.GetRandomString(7) 183 device_target = self.GetRandomString(4) 184 185 if lab_model: 186 lab = lab_model.name 187 elif lab_name: 188 lab = lab_name 189 else: 190 lab = self.GetRandomString() 191 192 schedule = model.ScheduleModel() 193 schedule.priority = priority 194 schedule.priority_value = Status.GetPriorityValue(schedule.priority) 195 schedule.period = period 196 schedule.shards = shards 197 schedule.retry_count = retry_count 198 schedule.required_signed_device_build = required_signed_device_build 199 schedule.build_storage_type = device_storage_type 200 schedule.manifest_branch = (device_branch if device_branch else 201 self.GetRandomString()) 202 schedule.build_target = "-".join([device_product, device_target]) 203 204 schedule.gsi_storage_type = gsi_storage_type 205 schedule.gsi_build_target = (gsi_build_target 206 if gsi_build_target else "-".join([ 207 self.GetRandomString(), 208 self.GetRandomString(4) 209 ])) 210 schedule.test_storage_type = test_storage_type 211 schedule.test_build_target = (test_build_target 212 if test_build_target else "-".join([ 213 self.GetRandomString(), 214 self.GetRandomString(4) 215 ])) 216 schedule.device = [] 217 schedule.device.append("/".join([lab, device_product])) 218 219 schedule.timestamp = datetime.datetime.now() 220 221 skip_list = [ 222 "priority", "priority_value", "period", "shards", 223 "retry_count", "required_signed_device_build", 224 "build_storage_type", "manifest_branch", "build_target", 225 "gsi_storage_type", "gsi_build_target", 226 "test_storage_type", "test_build_target", "device", 227 "children_jobs"] 228 set_or_empty = ["required_host_equipment", "required_device_equipment"] 229 for arg in schedule._properties: 230 if arg in skip_list or (arg in set_or_empty and arg not in kwargs): 231 continue 232 if arg in kwargs: 233 value = kwargs[arg] 234 elif isinstance(schedule._properties[arg], ndb.StringProperty): 235 value = self.GetRandomString() 236 elif isinstance(schedule._properties[arg], ndb.IntegerProperty): 237 value = 0 238 elif isinstance(schedule._properties[arg], ndb.BooleanProperty): 239 value = False 240 else: 241 print("A type of property '{}' is not supported.".format(arg)) 242 continue 243 if schedule._properties[arg]._repeated and type(value) is not list: 244 value = [value] 245 setattr(schedule, arg, value) 246 247 return schedule 248 249 def GenerateBuildModel(self, schedule, targets=None): 250 """Builds model.BuildModel with given information. 251 252 Args: 253 schedule: a model.ScheduleModel instance to look up build info. 254 targets: a list of strings which indicates artifact type. 255 256 Returns: 257 model.BuildModel instance. 258 """ 259 build_dict = {} 260 if targets is None: 261 targets = ["device", "gsi", "test"] 262 for target in targets: 263 build = model.BuildModel() 264 build.artifact_type = target 265 build.timestamp = datetime.datetime.now() 266 if target == "device": 267 build.signed = schedule.required_signed_device_build 268 build.manifest_branch = schedule.manifest_branch 269 build.build_target, build.build_type = ( 270 schedule.build_target.split("-")) 271 elif target == "gsi": 272 build.manifest_branch = schedule.gsi_branch 273 build.build_target, build.build_type = ( 274 schedule.gsi_build_target.split("-")) 275 elif target == "test": 276 build.manifest_branch = schedule.test_branch 277 build.build_target, build.build_type = ( 278 schedule.test_build_target.split("-")) 279 build.build_id = self.GetNewBuildId(build) 280 build_dict[target] = build 281 return build_dict 282 283 def GetNewBuildId(self, build): 284 """Generates build ID. 285 286 This method always generates newest (higher number) build ID than other 287 builds stored in testbed datastore. 288 289 Args: 290 build: a model.BuildModel instance to look up build information 291 from testbed datastore. 292 293 Returns: 294 a string, build ID. 295 """ 296 format_string = "{0:07d}" 297 build_query = model.BuildModel.query( 298 model.BuildModel.artifact_type == build.artifact_type, 299 model.BuildModel.build_target == build.build_target, 300 model.BuildModel.signed == build.signed, 301 model.BuildModel.manifest_branch == build.manifest_branch) 302 exiting_builds = build_query.fetch() 303 if exiting_builds: 304 exiting_builds.sort(key=lambda x: x.build_id, reverse=True) 305 latest_build_id = int(exiting_builds[0].build_id) 306 return format_string.format(latest_build_id + 1) 307 else: 308 return format_string.format(1) 309 310 def PassTime(self, hours=0, minutes=0, seconds=0): 311 """Assumes that a certain amount of time has passed. 312 313 This method changes does not change actual system time but changes all 314 jobs timestamp to assume time has passed. 315 316 Args: 317 hours: an integer, number of hours to pass time. 318 minutes: an integer, number of minutes to pass time. 319 seconds: an integer, number of seconds to pass time. 320 """ 321 if not hours and not minutes and not seconds: 322 return 323 324 jobs = model.JobModel.query().fetch() 325 to_put = [] 326 for job in jobs: 327 if job.timestamp: 328 job.timestamp -= datetime.timedelta( 329 hours=hours, minutes=minutes, seconds=seconds) 330 if job.heartbeat_stamp: 331 job.heartbeat_stamp -= datetime.timedelta( 332 hours=hours, minutes=minutes, seconds=seconds) 333 to_put.append(job) 334 if to_put: 335 ndb.put_multi(to_put) 336 337 def ResetDevices(self): 338 """Resets all devices to ready status.""" 339 devices = model.DeviceModel.query().fetch() 340 to_put = [] 341 for device in devices: 342 device.status = Status.DEVICE_STATUS_DICT["fastboot"] 343 device.scheduling_status = Status.DEVICE_SCHEDULING_STATUS_DICT[ 344 "free"] 345 to_put.append(device) 346 if to_put: 347 ndb.put_multi(to_put) 348