1 /* 2 * Copyright (C) 2017 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 package com.android.tradefed.invoker.shard; 17 18 import static org.junit.Assert.assertEquals; 19 import static org.junit.Assert.assertNotSame; 20 import static org.junit.Assert.assertTrue; 21 import static org.junit.Assert.fail; 22 import static org.mockito.Mockito.doReturn; 23 import static org.mockito.Mockito.doThrow; 24 25 import com.android.tradefed.build.BuildInfo; 26 import com.android.tradefed.build.IBuildInfo; 27 import com.android.tradefed.command.CommandOptions; 28 import com.android.tradefed.config.Configuration; 29 import com.android.tradefed.config.ConfigurationDef; 30 import com.android.tradefed.config.ConfigurationException; 31 import com.android.tradefed.config.ConfigurationFactory; 32 import com.android.tradefed.config.IConfiguration; 33 import com.android.tradefed.config.IGlobalConfiguration; 34 import com.android.tradefed.config.OptionSetter; 35 import com.android.tradefed.device.ITestDevice; 36 import com.android.tradefed.device.metric.BaseDeviceMetricCollector; 37 import com.android.tradefed.invoker.IInvocationContext; 38 import com.android.tradefed.invoker.IRescheduler; 39 import com.android.tradefed.invoker.InvocationContext; 40 import com.android.tradefed.result.ILogSaver; 41 import com.android.tradefed.suite.checker.KeyguardStatusChecker; 42 import com.android.tradefed.testtype.HostTest; 43 import com.android.tradefed.testtype.HostTestTest.SuccessTestCase; 44 import com.android.tradefed.testtype.HostTestTest.TestMetricTestCase; 45 import com.android.tradefed.testtype.IRemoteTest; 46 import com.android.tradefed.testtype.StubTest; 47 import com.android.tradefed.util.FileUtil; 48 import com.android.tradefed.util.keystore.IKeyStoreClient; 49 import com.android.tradefed.util.keystore.KeyStoreException; 50 import com.android.tradefed.util.keystore.StubKeyStoreFactory; 51 52 import org.junit.After; 53 import org.junit.Before; 54 import org.junit.Test; 55 import org.junit.runner.RunWith; 56 import org.junit.runners.JUnit4; 57 import org.mockito.ArgumentMatcher; 58 import org.mockito.Mockito; 59 60 import java.io.File; 61 import java.io.IOException; 62 import java.util.ArrayList; 63 import java.util.Arrays; 64 import java.util.List; 65 66 /** Unit tests for {@link ShardHelper}. */ 67 @RunWith(JUnit4.class) 68 public class ShardHelperTest { 69 70 private static final String TEST_CONFIG = 71 "<configuration description=\"shard config test\">\n" 72 + " <%s class=\"%s\" />\n" 73 + "</configuration>"; 74 75 private ShardHelper mHelper; 76 private IConfiguration mConfig; 77 private ILogSaver mMockLogSaver; 78 private IInvocationContext mContext; 79 private IRescheduler mRescheduler; 80 private IBuildInfo mBuildInfo; 81 82 @Before setUp()83 public void setUp() { 84 mHelper = 85 new ShardHelper() { 86 @Override 87 protected IGlobalConfiguration getGlobalConfiguration() { 88 try { 89 return ConfigurationFactory.getInstance() 90 .createGlobalConfigurationFromArgs( 91 new String[] {"empty"}, new ArrayList<>()); 92 } catch (ConfigurationException e) { 93 throw new RuntimeException(e); 94 } 95 } 96 97 @Override 98 protected void validateOptions(IConfiguration config) 99 throws ConfigurationException { 100 // Skip to avoid call to global configuration 101 } 102 }; 103 mConfig = new Configuration("fake_sharding_config", "desc"); 104 mContext = new InvocationContext(); 105 mBuildInfo = new BuildInfo(); 106 mContext.addDeviceBuildInfo(ConfigurationDef.DEFAULT_DEVICE_NAME, mBuildInfo); 107 mContext.addAllocatedDevice( 108 ConfigurationDef.DEFAULT_DEVICE_NAME, Mockito.mock(ITestDevice.class)); 109 mRescheduler = Mockito.mock(IRescheduler.class); 110 mMockLogSaver = Mockito.mock(ILogSaver.class); 111 mConfig.setLogSaver(mMockLogSaver); 112 } 113 114 @After tearDown()115 public void tearDown() { 116 mBuildInfo.cleanUp(); 117 } 118 119 /** 120 * Tests that when --shard-count is given to local sharding we create shard-count number of 121 * shards and not one shard per IRemoteTest. 122 */ 123 @Test testSplitWithShardCount()124 public void testSplitWithShardCount() throws Exception { 125 CommandOptions options = new CommandOptions(); 126 OptionSetter setter = new OptionSetter(options); 127 // shard-count is the number of shards we are requesting 128 setter.setOptionValue("shard-count", "3"); 129 mConfig.setCommandOptions(options); 130 mConfig.setCommandLine(new String[] {"empty"}); 131 StubTest test = new StubTest(); 132 setter = new OptionSetter(test); 133 // num-shards is a {@link StubTest} option that specify how many tests can stubtest split 134 // into. 135 setter.setOptionValue("num-shards", "5"); 136 mConfig.setTest(test); 137 assertEquals(1, mConfig.getTests().size()); 138 assertTrue(mHelper.shardConfig(mConfig, mContext, mRescheduler)); 139 // Ensure that we did split 1 tests per shard rescheduled. 140 Mockito.verify(mRescheduler, Mockito.times(3)) 141 .scheduleConfig( 142 Mockito.argThat( 143 new ArgumentMatcher<IConfiguration>() { 144 @Override 145 public boolean matches(IConfiguration argument) { 146 assertEquals(1, argument.getTests().size()); 147 return true; 148 } 149 })); 150 } 151 152 /** Tests that when --shard-count is not provided we create one shard per IRemoteTest. */ 153 @Test testSplit_noShardCount()154 public void testSplit_noShardCount() throws Exception { 155 CommandOptions options = new CommandOptions(); 156 OptionSetter setter = new OptionSetter(options); 157 mConfig.setCommandOptions(options); 158 mConfig.setCommandLine(new String[] {"empty"}); 159 StubTest test = new StubTest(); 160 setter = new OptionSetter(test); 161 setter.setOptionValue("num-shards", "5"); 162 mConfig.setTest(test); 163 assertEquals(1, mConfig.getTests().size()); 164 assertTrue(mHelper.shardConfig(mConfig, mContext, mRescheduler)); 165 // Ensure that we did split 1 tests per shard rescheduled. 166 Mockito.verify(mRescheduler, Mockito.times(5)) 167 .scheduleConfig( 168 Mockito.argThat( 169 new ArgumentMatcher<IConfiguration>() { 170 @Override 171 public boolean matches(IConfiguration argument) { 172 assertEquals(1, argument.getTests().size()); 173 return true; 174 } 175 })); 176 } 177 178 /** 179 * Tests that when a --shard-count 10 is requested but there is only 5 sub tests after sharding 180 * there is no point in rescheduling 10 times so we limit to the number of tests. 181 */ 182 @Test testSplitWithShardCount_notEnoughTest()183 public void testSplitWithShardCount_notEnoughTest() throws Exception { 184 CommandOptions options = new CommandOptions(); 185 OptionSetter setter = new OptionSetter(options); 186 setter.setOptionValue("shard-count", "10"); 187 mConfig.setCommandOptions(options); 188 mConfig.setCommandLine(new String[] {"empty"}); 189 StubTest test = new StubTest(); 190 setter = new OptionSetter(test); 191 // num-shards is a {@link StubTest} option that specify how many tests can stubtest split 192 // into. 193 setter.setOptionValue("num-shards", "5"); 194 mConfig.setTest(test); 195 assertEquals(1, mConfig.getTests().size()); 196 assertTrue(mHelper.shardConfig(mConfig, mContext, mRescheduler)); 197 // We only reschedule 5 times and not 10 like --shard-count because there is not enough 198 // tests to put at least 1 test per shard. So there is no point in rescheduling on new 199 // devices. 200 Mockito.verify(mRescheduler, Mockito.times(5)) 201 .scheduleConfig( 202 Mockito.argThat( 203 new ArgumentMatcher<IConfiguration>() { 204 @Override 205 public boolean matches(IConfiguration argument) { 206 assertEquals(1, argument.getTests().size()); 207 return true; 208 } 209 })); 210 } 211 createTmpConfig(String objType, Object obj)212 private File createTmpConfig(String objType, Object obj) throws IOException { 213 File configFile = FileUtil.createTempFile("shard-helper-test", ".xml"); 214 String content = String.format(TEST_CONFIG, objType, obj.getClass().getCanonicalName()); 215 FileUtil.writeToFile(content, configFile); 216 return configFile; 217 } 218 219 /** 220 * Test that some objects are being cloned to the shards in order to avoid shared state issues. 221 */ 222 @Test testCloneStatusChecker()223 public void testCloneStatusChecker() throws Exception { 224 KeyguardStatusChecker checker = new KeyguardStatusChecker(); 225 File configFile = createTmpConfig(Configuration.SYSTEM_STATUS_CHECKER_TYPE_NAME, checker); 226 try { 227 CommandOptions options = new CommandOptions(); 228 OptionSetter setter = new OptionSetter(options); 229 // shard-count is the number of shards we are requesting 230 setter.setOptionValue("shard-count", "3"); 231 mConfig.setCommandOptions(options); 232 mConfig.setCommandLine(new String[] {configFile.getAbsolutePath()}); 233 mConfig.setSystemStatusChecker(checker); 234 StubTest test = new StubTest(); 235 setter = new OptionSetter(test); 236 // num-shards is a {@link StubTest} option that specify how many tests can stubtest split 237 // into. 238 setter.setOptionValue("num-shards", "5"); 239 mConfig.setTest(test); 240 assertEquals(1, mConfig.getTests().size()); 241 assertTrue(mHelper.shardConfig(mConfig, mContext, mRescheduler)); 242 // Ensure that we did split 1 tests per shard rescheduled. 243 Mockito.verify(mRescheduler, Mockito.times(3)) 244 .scheduleConfig( 245 Mockito.argThat( 246 new ArgumentMatcher<IConfiguration>() { 247 @Override 248 public boolean matches(IConfiguration argument) { 249 assertEquals(1, argument.getTests().size()); 250 // Status checker is in all shard and a new one has been 251 // created 252 assertEquals( 253 1, argument.getSystemStatusCheckers().size()); 254 assertNotSame( 255 checker, 256 argument.getSystemStatusCheckers().get(0)); 257 return true; 258 } 259 })); 260 } finally { 261 FileUtil.deleteFile(configFile); 262 } 263 } 264 265 /** 266 * Test that some objects are being cloned to the shards in order to avoid shared state issues. 267 */ 268 @Test testCloneMetricCollector()269 public void testCloneMetricCollector() throws Exception { 270 BaseDeviceMetricCollector collector = new BaseDeviceMetricCollector(); 271 File configFile = 272 createTmpConfig(Configuration.DEVICE_METRICS_COLLECTOR_TYPE_NAME, collector); 273 try { 274 CommandOptions options = new CommandOptions(); 275 OptionSetter setter = new OptionSetter(options); 276 // shard-count is the number of shards we are requesting 277 setter.setOptionValue("shard-count", "3"); 278 mConfig.setCommandOptions(options); 279 mConfig.setCommandLine(new String[] {configFile.getAbsolutePath()}); 280 mConfig.setDeviceMetricCollectors(Arrays.asList(collector)); 281 StubTest test = new StubTest(); 282 setter = new OptionSetter(test); 283 // num-shards is a {@link StubTest} option that specify how many tests can stubtest split 284 // into. 285 setter.setOptionValue("num-shards", "5"); 286 mConfig.setTest(test); 287 assertEquals(1, mConfig.getTests().size()); 288 assertTrue(mHelper.shardConfig(mConfig, mContext, mRescheduler)); 289 // Ensure that we did split 1 tests per shard rescheduled. 290 Mockito.verify(mRescheduler, Mockito.times(3)) 291 .scheduleConfig( 292 Mockito.argThat( 293 new ArgumentMatcher<IConfiguration>() { 294 @Override 295 public boolean matches(IConfiguration argument) { 296 assertEquals(1, argument.getTests().size()); 297 // Status checker is in all shard and a new one has been 298 // created 299 assertEquals(1, argument.getMetricCollectors().size()); 300 assertNotSame( 301 collector, 302 argument.getMetricCollectors().get(0)); 303 return true; 304 } 305 })); 306 } finally { 307 FileUtil.deleteFile(configFile); 308 } 309 } 310 311 /** 312 * Test that even when sharding, configuration are loaded with the global keystore if needed. 313 */ 314 @Test testClone_withKeystore()315 public void testClone_withKeystore() throws Exception { 316 IKeyStoreClient mockClient = Mockito.mock(IKeyStoreClient.class); 317 mHelper = 318 new ShardHelper() { 319 @Override 320 protected IGlobalConfiguration getGlobalConfiguration() { 321 try { 322 IGlobalConfiguration config = 323 ConfigurationFactory.getInstance() 324 .createGlobalConfigurationFromArgs( 325 new String[] {"empty"}, new ArrayList<>()); 326 config.setKeyStoreFactory( 327 new StubKeyStoreFactory() { 328 @Override 329 public IKeyStoreClient createKeyStoreClient() 330 throws KeyStoreException { 331 return mockClient; 332 } 333 }); 334 return config; 335 } catch (ConfigurationException e) { 336 throw new RuntimeException(e); 337 } 338 } 339 340 @Override 341 protected void validateOptions(IConfiguration config) 342 throws ConfigurationException { 343 // Skip to avoid call to global configuration 344 } 345 }; 346 CommandOptions options = new CommandOptions(); 347 HostTest stubTest = new HostTest(); 348 OptionSetter setter = new OptionSetter(options, stubTest); 349 // shard-count is the number of shards we are requesting 350 setter.setOptionValue("shard-count", "2"); 351 setter.setOptionValue("class", SuccessTestCase.class.getName()); 352 setter.setOptionValue("class", TestMetricTestCase.class.getName()); 353 mConfig.setCommandOptions(options); 354 mConfig.setCommandLine(new String[] {"host", "--class", "USE_KEYSTORE@test"}); 355 mConfig.setTest(stubTest); 356 357 assertEquals(1, mConfig.getTests().size()); 358 359 doReturn(true).when(mockClient).isAvailable(); 360 doReturn(SuccessTestCase.class.getName()).when(mockClient).fetchKey("test"); 361 362 assertTrue(mHelper.shardConfig(mConfig, mContext, mRescheduler)); 363 // Ensure that we did split 1 tests per shard rescheduled. 364 Mockito.verify(mRescheduler, Mockito.times(2)) 365 .scheduleConfig( 366 Mockito.argThat( 367 new ArgumentMatcher<IConfiguration>() { 368 @Override 369 public boolean matches(IConfiguration argument) { 370 assertEquals(1, argument.getTests().size()); 371 return true; 372 } 373 })); 374 } 375 376 /** 377 * Test that even when sharding, configuration are loaded with the global keystore if needed. If 378 * the keystore fails to load a key, the invocation will be stopped. 379 */ 380 @Test testClone_withKeystore_loadingFails()381 public void testClone_withKeystore_loadingFails() throws Exception { 382 IKeyStoreClient mockClient = Mockito.mock(IKeyStoreClient.class); 383 mHelper = 384 new ShardHelper() { 385 @Override 386 protected IGlobalConfiguration getGlobalConfiguration() { 387 try { 388 IGlobalConfiguration config = 389 ConfigurationFactory.getInstance() 390 .createGlobalConfigurationFromArgs( 391 new String[] {"empty"}, new ArrayList<>()); 392 config.setKeyStoreFactory( 393 new StubKeyStoreFactory() { 394 @Override 395 public IKeyStoreClient createKeyStoreClient() 396 throws KeyStoreException { 397 return mockClient; 398 } 399 }); 400 return config; 401 } catch (ConfigurationException e) { 402 throw new RuntimeException(e); 403 } 404 } 405 }; 406 CommandOptions options = new CommandOptions(); 407 HostTest stubTest = new HostTest(); 408 OptionSetter setter = new OptionSetter(options, stubTest); 409 // shard-count is the number of shards we are requesting 410 setter.setOptionValue("shard-count", "2"); 411 setter.setOptionValue("class", SuccessTestCase.class.getName()); 412 setter.setOptionValue("class", TestMetricTestCase.class.getName()); 413 mConfig.setCommandOptions(options); 414 mConfig.setCommandLine(new String[] {"host", "--class", "USE_KEYSTORE@test"}); 415 mConfig.setTest(stubTest); 416 417 assertEquals(1, mConfig.getTests().size()); 418 419 doReturn(true).when(mockClient).isAvailable(); 420 doReturn(SuccessTestCase.class.getName()).when(mockClient).fetchKey("test"); 421 doThrow(new RuntimeException()).when(mockClient).fetchKey("test"); 422 try { 423 mHelper.shardConfig(mConfig, mContext, mRescheduler); 424 fail("Should have thrown an exception."); 425 } catch (RuntimeException expected) { 426 // expected 427 } 428 } 429 430 /** 431 * Test split when a token test is present. The token pool of each shard should be populated. 432 */ 433 @Test testSplitWithTokens()434 public void testSplitWithTokens() throws Exception { 435 CommandOptions options = new CommandOptions(); 436 OptionSetter setter = new OptionSetter(options); 437 // shard-count is the number of shards we are requesting 438 setter.setOptionValue("shard-count", "3"); 439 // Enable token sharding 440 setter.setOptionValue("enable-token-sharding", "true"); 441 442 mConfig.setCommandOptions(options); 443 mConfig.setCommandLine(new String[] {"empty"}); 444 List<IRemoteTest> tests = new ArrayList<>(); 445 StubTest test = new StubTest(); 446 setter = new OptionSetter(test); 447 // num-shards is a {@link StubTest} option that specify how many tests can stubtest split 448 // into. 449 setter.setOptionValue("num-shards", "5"); 450 tests.add(test); 451 // Add a token test 452 tests.add(new TokenTestClass()); 453 mConfig.setTests(tests); 454 assertEquals(2, mConfig.getTests().size()); 455 assertTrue(mHelper.shardConfig(mConfig, mContext, mRescheduler)); 456 // Ensure that we did split 1 tests per shard rescheduled. 457 Mockito.verify(mRescheduler, Mockito.times(3)) 458 .scheduleConfig( 459 Mockito.argThat( 460 new ArgumentMatcher<IConfiguration>() { 461 @Override 462 public boolean matches(IConfiguration argument) { 463 assertEquals(1, argument.getTests().size()); 464 IRemoteTest test = argument.getTests().get(0); 465 assertTrue(test instanceof TestsPoolPoller); 466 TestsPoolPoller poller = (TestsPoolPoller) test; 467 // Token pool has the test 468 assertEquals(1, poller.getTokenPool().size()); 469 return true; 470 } 471 })); 472 } 473 } 474