1 /*
2  * Copyright (C) 2016 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.testtype;
17 
18 import static com.google.common.truth.Truth.assertThat;
19 
20 import static org.mockito.AdditionalMatchers.gt;
21 import static org.mockito.ArgumentMatchers.any;
22 import static org.mockito.ArgumentMatchers.anyCollection;
23 import static org.mockito.ArgumentMatchers.anyInt;
24 import static org.mockito.ArgumentMatchers.anyString;
25 import static org.mockito.ArgumentMatchers.eq;
26 import static org.mockito.Mockito.doAnswer;
27 import static org.mockito.Mockito.doReturn;
28 import static org.mockito.Mockito.mock;
29 import static org.mockito.Mockito.verify;
30 
31 import com.android.ddmlib.IDevice;
32 import com.android.ddmlib.testrunner.IRemoteAndroidTestRunner;
33 import com.android.ddmlib.testrunner.ITestRunListener;
34 import com.android.ddmlib.testrunner.RemoteAndroidTestRunner;
35 import com.android.tradefed.build.BuildInfo;
36 import com.android.tradefed.build.IBuildInfo;
37 import com.android.tradefed.device.DeviceNotAvailableException;
38 import com.android.tradefed.device.ITestDevice;
39 import com.android.tradefed.log.ITestLogger;
40 import com.android.tradefed.result.ITestInvocationListener;
41 import com.android.tradefed.result.ITestLifeCycleReceiver;
42 import com.android.tradefed.result.InputStreamSource;
43 import com.android.tradefed.result.LogDataType;
44 import com.android.tradefed.result.TestDescription;
45 import com.android.tradefed.result.TestRunResult;
46 import com.android.tradefed.result.ddmlib.TestRunToTestInvocationForwarder;
47 import com.android.tradefed.util.ICompressionStrategy;
48 import com.android.tradefed.util.ListInstrumentationParser;
49 import com.android.tradefed.util.ListInstrumentationParser.InstrumentationTarget;
50 
51 import com.google.common.collect.ImmutableList;
52 import com.google.common.collect.ImmutableMap;
53 import com.google.common.collect.Lists;
54 import com.google.protobuf.ByteString;
55 
56 import org.junit.Before;
57 import org.junit.Rule;
58 import org.junit.Test;
59 import org.junit.rules.TemporaryFolder;
60 import org.junit.runner.RunWith;
61 import org.junit.runners.JUnit4;
62 import org.mockito.Mock;
63 import org.mockito.Mockito;
64 import org.mockito.MockitoAnnotations;
65 import org.mockito.invocation.InvocationOnMock;
66 import org.mockito.stubbing.Answer;
67 
68 import java.io.File;
69 import java.io.FileInputStream;
70 import java.io.FileOutputStream;
71 import java.io.IOException;
72 import java.io.InputStream;
73 import java.util.ArrayList;
74 import java.util.Collection;
75 import java.util.HashMap;
76 import java.util.List;
77 import java.util.Map;
78 
79 /** Unit tests for {@link CodeCoverageTestBase}. */
80 @RunWith(JUnit4.class)
81 public class CodeCoverageTestBaseTest {
82 
83     private static final long TEST_RUN_TIME = 1000;
84     private static final String COVERAGE_PATH = "/data/user/0/%s/files/coverage.ec";
85 
86     private static final String PACKAGE_NAME1 = "com.example.foo.test";
87     private static final String PACKAGE_NAME2 = "com.example.bar.test";
88     private static final String PACKAGE_NAME3 = "com.example.baz.test";
89 
90     private static final String RUNNER_NAME1 = "android.support.test.runner.AndroidJUnitRunner";
91     private static final String RUNNER_NAME2 = "android.test.InstrumentationTestRunner";
92     private static final String RUNNER_NAME3 = "com.example.custom.Runner";
93 
94     private static final TestDescription FOO_TEST1 = new TestDescription(".FooTest", "test1");
95     private static final TestDescription FOO_TEST2 = new TestDescription(".FooTest", "test2");
96     private static final TestDescription FOO_TEST3 = new TestDescription(".FooTest", "test3");
97     private static final List<TestDescription> FOO_TESTS =
98             ImmutableList.of(FOO_TEST1, FOO_TEST2, FOO_TEST3);
99 
100     private static final TestDescription BAR_TEST1 = new TestDescription(".BarTest", "test1");
101     private static final TestDescription BAR_TEST2 = new TestDescription(".BarTest", "test2");
102     private static final List<TestDescription> BAR_TESTS = ImmutableList.of(BAR_TEST1, BAR_TEST2);
103 
104     private static final TestDescription BAZ_TEST1 = new TestDescription(".BazTest", "test1");
105     private static final TestDescription BAZ_TEST2 = new TestDescription(".BazTest", "test2");
106     private static final List<TestDescription> BAZ_TESTS = ImmutableList.of(BAZ_TEST1, BAZ_TEST2);
107 
108     private static final ByteString FAKE_REPORT_CONTENTS =
109             ByteString.copyFromUtf8("Mi estas kovrado raporto");
110 
111     private static final ByteString FAKE_MEASUREMENT1 =
112             ByteString.copyFromUtf8("Mi estas kovrado mezurado");
113     private static final ByteString FAKE_MEASUREMENT2 =
114             ByteString.copyFromUtf8("Mi estas ankau kovrado mezurado");
115     private static final ByteString FAKE_MEASUREMENT3 =
116             ByteString.copyFromUtf8("Mi estas ankorau alia priraportado mezurado");
117 
118     private static final IBuildInfo BUILD_INFO = new BuildInfo("123456", "device-userdebug");
119 
120     @Rule public TemporaryFolder mFolder = new TemporaryFolder();
121 
122     // Mocks
123     @Mock ITestDevice mDevice;
124 
125     @Mock ITestInvocationListener mListener;
126 
127     @Mock ListInstrumentationParser mInstrumentationParser;
128 
129     // Fake test data
130     @Mock TestDataRegistry<List<TestDescription>> mTests;
131 
132     @Mock TestDataRegistry<ByteString> mMeasurements;
133 
134     interface TestDataRegistry<T> {
get(String packageName, String runnerName, int shardIndex, int numShards)135         T get(String packageName, String runnerName, int shardIndex, int numShards);
136     }
137 
138     /** Object under test */
139     CodeCoverageTestStub mCoverageTest;
140 
141     @Before
setUp()142     public void setUp() throws DeviceNotAvailableException {
143         MockitoAnnotations.initMocks(this);
144         doAnswer(CALL_RUNNER)
145                 .when(mDevice)
146                 .runInstrumentationTests(
147                         any(IRemoteAndroidTestRunner.class), any(ITestLifeCycleReceiver.class));
148         mCoverageTest = new CodeCoverageTestStub();
149     }
150 
151     static enum FakeReportFormat implements CodeCoverageReportFormat {
152         CSV(LogDataType.JACOCO_CSV),
153         XML(LogDataType.JACOCO_XML),
154         HTML(LogDataType.HTML);
155 
156         private final LogDataType mLogDataType;
157 
FakeReportFormat(LogDataType logDataType)158         private FakeReportFormat(LogDataType logDataType) {
159             mLogDataType = logDataType;
160         }
161 
162         @Override
getLogDataType()163         public LogDataType getLogDataType() { return mLogDataType; }
164     }
165 
166     /** A subclass of {@link CodeCoverageTest} with certain methods stubbed out for testing. */
167     private class CodeCoverageTestStub extends CodeCoverageTestBase<FakeReportFormat> {
168 
169         // Captured data
170         private ImmutableList.Builder<ByteString> mCapturedMeasurements =
171                 new ImmutableList.Builder<>();
172 
CodeCoverageTestStub()173         public CodeCoverageTestStub() {
174             setDevice(mDevice);
175         }
176 
177         @Override
getBuild()178         public IBuildInfo getBuild() { return BUILD_INFO; }
179 
180         @Override
getCompressionStrategy()181         ICompressionStrategy getCompressionStrategy() {
182             return mock(ICompressionStrategy.class);
183         }
184 
getMeasurements()185         ImmutableList<ByteString> getMeasurements() {
186             return mCapturedMeasurements.build();
187         }
188 
189         @Override
generateCoverageReport( Collection<File> measurementFiles, FakeReportFormat format)190         protected File generateCoverageReport(
191                 Collection<File> measurementFiles, FakeReportFormat format) throws IOException {
192             // Capture the measurements for verification later
193             for (File measurementFile : measurementFiles) {
194                 try (FileInputStream inputStream = new FileInputStream(measurementFile)) {
195                     mCapturedMeasurements.add(ByteString.readFrom(inputStream));
196                 }
197             }
198 
199             // Write the fake report
200             File ret = mFolder.newFile();
201             FAKE_REPORT_CONTENTS.writeTo(new FileOutputStream(ret));
202             return ret;
203         }
204 
205         @Override
getReportFormat()206         protected List<FakeReportFormat> getReportFormat() {
207             return ImmutableList.of(FakeReportFormat.HTML);
208         }
209 
210         @Override
internalCreateTest()211         InstrumentationTest internalCreateTest() {
212             return new InstrumentationTest() {
213                 @Override
214                 IRemoteAndroidTestRunner createRemoteAndroidTestRunner(
215                         String packageName, String runnerName, IDevice device) {
216                     return new FakeTestRunner(packageName, runnerName);
217                 }
218             };
219         }
220 
221         @Override
internalCreateTestRunner(String packageName, String runnerName)222         IRemoteAndroidTestRunner internalCreateTestRunner(String packageName, String runnerName) {
223             return new FakeTestRunner(packageName, runnerName);
224         }
225 
226         @Override
internalCreateListInstrumentationParser()227         ListInstrumentationParser internalCreateListInstrumentationParser() {
228             return mInstrumentationParser;
229         }
230     }
231 
232     private static final class LogCaptorAnswer implements Answer<Void> {
233         private ByteString mValue;
234 
235         @Override
236         public Void answer(InvocationOnMock invocation) throws Throwable {
237             Object[] args = invocation.getArguments();
238             InputStream reportStream = ((InputStreamSource) args[2]).createInputStream();
239             mValue = ByteString.readFrom(reportStream);
240             return null;
241         }
242 
243         ByteString getValue() {
244             return mValue;
245         }
246     }
247 
248     @Test
249     public void testRun() throws DeviceNotAvailableException {
250         // Prepare some test data
251         doReturn(ImmutableList.of(new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "")))
252                 .when(mInstrumentationParser)
253                 .getInstrumentationTargets();
254 
255         // Mocking boilerplate
256         doReturn(FOO_TESTS)
257                 .when(mTests)
258                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
259         doReturn(FAKE_MEASUREMENT1)
260                 .when(mMeasurements)
261                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
262 
263         // Validate the report when it gets log, since the file will get cleaned up later
264         LogCaptorAnswer logCaptor = new LogCaptorAnswer();
265         doAnswer(logCaptor)
266                 .when(mListener)
267                 .testLog(
268                         eq("coverage"),
269                         eq(FakeReportFormat.HTML.getLogDataType()),
270                         any(InputStreamSource.class));
271 
272         // Run the test
273         mCoverageTest.run(mListener);
274 
275         // Verify that the measurements were collected and the report was logged
276         assertThat(mCoverageTest.getMeasurements()).containsExactly(FAKE_MEASUREMENT1);
277         assertThat(logCaptor.getValue()).isEqualTo(FAKE_REPORT_CONTENTS);
278     }
279 
280     @Test
281     public void testRun_multipleInstrumentationTargets() throws DeviceNotAvailableException {
282         // Prepare some test data
283         ImmutableList<InstrumentationTarget> targets =
284                 ImmutableList.of(
285                         new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, ""),
286                         new InstrumentationTarget(PACKAGE_NAME2, RUNNER_NAME1, ""),
287                         new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME1, ""));
288         doReturn(targets).when(mInstrumentationParser).getInstrumentationTargets();
289 
290         // Mocking boilerplate
291         doReturn(FOO_TESTS)
292                 .when(mTests)
293                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
294         doReturn(FAKE_MEASUREMENT1)
295                 .when(mMeasurements)
296                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
297 
298         doReturn(BAR_TESTS)
299                 .when(mTests)
300                 .get(eq(PACKAGE_NAME2), eq(RUNNER_NAME1), anyInt(), anyInt());
301         doReturn(FAKE_MEASUREMENT2)
302                 .when(mMeasurements)
303                 .get(eq(PACKAGE_NAME2), eq(RUNNER_NAME1), anyInt(), anyInt());
304 
305         doReturn(BAZ_TESTS)
306                 .when(mTests)
307                 .get(eq(PACKAGE_NAME3), eq(RUNNER_NAME1), anyInt(), anyInt());
308         doReturn(FAKE_MEASUREMENT3)
309                 .when(mMeasurements)
310                 .get(eq(PACKAGE_NAME3), eq(RUNNER_NAME1), anyInt(), anyInt());
311 
312         // Run the test
313         mCoverageTest.run(mListener);
314 
315         // Verify that all targets were run by checking that we recieved measurements from each
316         assertThat(mCoverageTest.getMeasurements())
317                 .containsExactly(FAKE_MEASUREMENT1, FAKE_MEASUREMENT2, FAKE_MEASUREMENT3);
318     }
319 
320     @Test
321     public void testRun_multipleShards() throws DeviceNotAvailableException {
322         // Prepare some test data
323         doReturn(ImmutableList.of(new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "")))
324                 .when(mInstrumentationParser)
325                 .getInstrumentationTargets();
326         List<List<TestDescription>> shards = Lists.partition(FOO_TESTS, 1);
327 
328         // Indicate that the test should be split into 3 shards
329         doReturn(FOO_TESTS).when(mTests).get(PACKAGE_NAME1, RUNNER_NAME1, 0, 1);
330         doReturn(FOO_TESTS.subList(0, FOO_TESTS.size() / 2))
331                 .when(mTests)
332                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(0), eq(2));
333         mCoverageTest.setMaxTestsPerChunk(1);
334 
335         // Return subsets of FOO_TESTS when running shards
336         doReturn(shards.get(0)).when(mTests).get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(0), gt(1));
337         doReturn(FAKE_MEASUREMENT1)
338                 .when(mMeasurements)
339                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(0), gt(1));
340 
341         doReturn(shards.get(1)).when(mTests).get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(1), gt(1));
342         doReturn(FAKE_MEASUREMENT2)
343                 .when(mMeasurements)
344                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(1), gt(1));
345 
346         doReturn(shards.get(2)).when(mTests).get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(2), gt(1));
347         doReturn(FAKE_MEASUREMENT3)
348                 .when(mMeasurements)
349                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(2), gt(1));
350 
351         // Run the test
352         mCoverageTest.run(mListener);
353 
354         // Verify that all shards were run by checking that we recieved measurements from each
355         assertThat(mCoverageTest.getMeasurements())
356                 .containsExactly(FAKE_MEASUREMENT1, FAKE_MEASUREMENT2, FAKE_MEASUREMENT3);
357     }
358 
359     @Test
360     public void testRun_rerunIndividualTests_failedRun() throws DeviceNotAvailableException {
361         // Prepare some test data
362         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
363         doReturn(ImmutableList.of(target)).when(mInstrumentationParser).getInstrumentationTargets();
364         doReturn(FOO_TESTS)
365                 .when(mTests)
366                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
367 
368         TestRunResult success = mock(TestRunResult.class);
369         doReturn(false).when(success).isRunFailure();
370         TestRunResult failure = mock(TestRunResult.class);
371         doReturn(true).when(failure).isRunFailure();
372 
373         // Mocking boilerplate
374         ITestInvocationListener mockListener = mock(ITestInvocationListener.class);
375         CodeCoverageTestStub coverageTest = Mockito.spy(new CodeCoverageTestStub());
376         doReturn(FOO_TESTS).when(mTests).get(PACKAGE_NAME1, RUNNER_NAME1, 0, 1);
377 
378         doReturn(failure).when(coverageTest).runTest(eq(target), eq(0), eq(1),
379                 any(ITestInvocationListener.class));
380         doReturn(success).when(coverageTest).runTest(eq(target), eq(FOO_TEST1),
381                 any(ITestInvocationListener.class));
382         doReturn(failure).when(coverageTest).runTest(eq(target), eq(FOO_TEST2),
383                 any(ITestInvocationListener.class));
384         doReturn(success).when(coverageTest).runTest(eq(target), eq(FOO_TEST3),
385                 any(ITestInvocationListener.class));
386 
387         // Run the test
388         coverageTest.run(mockListener);
389 
390         // Verify that individual tests are rerun
391         verify(coverageTest).runTest(eq(target), eq(0), eq(1), any(ITestInvocationListener.class));
392         verify(coverageTest).runTest(eq(target), eq(FOO_TEST1), any(ITestInvocationListener.class));
393         verify(coverageTest).runTest(eq(target), eq(FOO_TEST2), any(ITestInvocationListener.class));
394         verify(coverageTest).runTest(eq(target), eq(FOO_TEST3), any(ITestInvocationListener.class));
395     }
396 
397     @Test
398     public void testRun_rerunIndividualTests_missingCoverageFile()
399             throws DeviceNotAvailableException {
400         // Prepare some test data
401         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
402         doReturn(ImmutableList.of(target)).when(mInstrumentationParser).getInstrumentationTargets();
403         doReturn(FOO_TESTS)
404                 .when(mTests)
405                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
406 
407         TestRunResult success = mock(TestRunResult.class);
408         doReturn(false).when(success).isRunFailure();
409 
410         // Mocking boilerplate
411         ITestInvocationListener mockListener = mock(ITestInvocationListener.class);
412         CodeCoverageTestStub coverageTest = Mockito.spy(new CodeCoverageTestStub());
413         doReturn(FOO_TESTS).when(mTests).get(PACKAGE_NAME1, RUNNER_NAME1, 0, 1);
414 
415         doReturn(success)
416                 .when(coverageTest)
417                 .runTest(any(InstrumentationTest.class), any(ITestInvocationListener.class));
418         ITestDevice mockDevice = coverageTest.getDevice();
419         doReturn(false).doReturn(true).when(mockDevice).doesFileExist(anyString());
420 
421         // Run the test
422         coverageTest.run(mockListener);
423 
424         // Verify that individual tests are rerun
425         verify(coverageTest).runTest(eq(target), eq(0), eq(1), any(ITestInvocationListener.class));
426         verify(coverageTest).runTest(eq(target), eq(FOO_TEST1), any(ITestInvocationListener.class));
427         verify(coverageTest).runTest(eq(target), eq(FOO_TEST2), any(ITestInvocationListener.class));
428         verify(coverageTest).runTest(eq(target), eq(FOO_TEST3), any(ITestInvocationListener.class));
429     }
430 
431     @Test
432     public void testRun_multipleFormats() throws DeviceNotAvailableException, IOException {
433         // Prepare some test data
434         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
435         doReturn(ImmutableList.of(target)).when(mInstrumentationParser).getInstrumentationTargets();
436         doReturn(FOO_TESTS)
437                 .when(mTests)
438                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
439         doReturn(FAKE_MEASUREMENT1)
440                 .when(mMeasurements)
441                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
442 
443         File fakeHtmlReport = new File("/some/fake/xml/report/");
444         File fakeXmlReport = new File("/some/fake/xml/report.xml");
445 
446         // Mocking boilerplate
447         ITestInvocationListener mockListener = mock(ITestInvocationListener.class);
448         CodeCoverageTestStub coverageTest = Mockito.spy(new CodeCoverageTestStub());
449         doReturn(FOO_TESTS).when(mTests).get(PACKAGE_NAME1, RUNNER_NAME1, 0, 1);
450         doReturn(ImmutableList.of(FakeReportFormat.XML, FakeReportFormat.HTML))
451                 .when(coverageTest)
452                 .getReportFormat();
453         doReturn(fakeHtmlReport)
454                 .when(coverageTest)
455                 .generateCoverageReport(anyCollection(), eq(FakeReportFormat.HTML));
456         doReturn(fakeXmlReport)
457                 .when(coverageTest)
458                 .generateCoverageReport(anyCollection(), eq(FakeReportFormat.XML));
459 
460         // Run the test
461         coverageTest.run(mockListener);
462 
463         // Verify that the test was run, and that the reports were logged
464         verify(coverageTest)
465                 .runTest(any(InstrumentationTest.class), any(ITestInvocationListener.class));
466         verify(coverageTest).generateCoverageReport(anyCollection(), eq(FakeReportFormat.HTML));
467         verify(coverageTest).doLogReport(anyString(), eq(FakeReportFormat.HTML.getLogDataType()),
468                 eq(fakeHtmlReport), any(ITestLogger.class));
469         verify(coverageTest).generateCoverageReport(anyCollection(), eq(FakeReportFormat.XML));
470         verify(coverageTest).doLogReport(anyString(), eq(FakeReportFormat.XML.getLogDataType()),
471                 eq(fakeXmlReport), any(ITestLogger.class));
472     }
473 
474     @Test
475     public void testGetInstrumentationTargets() throws DeviceNotAvailableException {
476         // Prepare some test data
477         InstrumentationTarget target1 = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
478         InstrumentationTarget target2 = new InstrumentationTarget(PACKAGE_NAME2, RUNNER_NAME2, "");
479         InstrumentationTarget target3 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME1, "");
480         InstrumentationTarget target4 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME3, "");
481 
482         // Set up mocks
483         doReturn(ImmutableList.of(target1, target2, target3, target4))
484                 .when(mInstrumentationParser)
485                 .getInstrumentationTargets();
486 
487         // Get the instrumentation targets
488         Collection<InstrumentationTarget> targets = mCoverageTest.getInstrumentationTargets();
489 
490         // Verify that all of the instrumentation targets were found
491         assertThat(targets).containsExactly(target1, target2, target3, target4);
492     }
493 
494     @Test
495     public void testGetInstrumentationTargets_packageFilterSingleFilterSingleResult()
496             throws DeviceNotAvailableException {
497         // Prepare some test data
498         InstrumentationTarget target1 = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
499         InstrumentationTarget target2 = new InstrumentationTarget(PACKAGE_NAME2, RUNNER_NAME2, "");
500         InstrumentationTarget target3 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME1, "");
501         InstrumentationTarget target4 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME3, "");
502 
503         // Set up mocks
504         doReturn(ImmutableList.of(target1, target2, target3, target4))
505                 .when(mInstrumentationParser)
506                 .getInstrumentationTargets();
507         mCoverageTest.setPackageFilter(ImmutableList.of(PACKAGE_NAME1));
508 
509         // Get the instrumentation targets
510         Collection<InstrumentationTarget> targets = mCoverageTest.getInstrumentationTargets();
511 
512         // Verify that only the PACKAGE_NAME1 target was returned
513         assertThat(targets).containsExactly(target1);
514     }
515 
516     @Test
517     public void testGetInstrumentationTargets_packageFilterSingleFilterMultipleResults()
518             throws DeviceNotAvailableException {
519         // Prepare some test data
520         InstrumentationTarget target1 = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
521         InstrumentationTarget target2 = new InstrumentationTarget(PACKAGE_NAME2, RUNNER_NAME2, "");
522         InstrumentationTarget target3 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME1, "");
523         InstrumentationTarget target4 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME3, "");
524 
525         // Set up mocks
526         doReturn(ImmutableList.of(target1, target2, target3, target4))
527                 .when(mInstrumentationParser)
528                 .getInstrumentationTargets();
529         mCoverageTest.setPackageFilter(ImmutableList.of(PACKAGE_NAME3));
530 
531         // Get the instrumentation targets
532         Collection<InstrumentationTarget> targets = mCoverageTest.getInstrumentationTargets();
533 
534         // Verify that both PACKAGE_NAME3 targets were returned
535         assertThat(targets).containsExactly(target3, target4);
536     }
537 
538     @Test
539     public void testGetInstrumentationTargets_packageFilterMultipleFilters()
540             throws DeviceNotAvailableException {
541         // Prepare some test data
542         InstrumentationTarget target1 = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
543         InstrumentationTarget target2 = new InstrumentationTarget(PACKAGE_NAME2, RUNNER_NAME2, "");
544         InstrumentationTarget target3 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME1, "");
545         InstrumentationTarget target4 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME3, "");
546 
547         // Set up mocks
548         doReturn(ImmutableList.of(target1, target2, target3, target4))
549                 .when(mInstrumentationParser)
550                 .getInstrumentationTargets();
551         mCoverageTest.setPackageFilter(ImmutableList.of(PACKAGE_NAME1, PACKAGE_NAME2));
552 
553         // Get the instrumentation targets
554         Collection<InstrumentationTarget> targets = mCoverageTest.getInstrumentationTargets();
555 
556         // Verify that the PACKAGE_NAME1 and PACKAGE_NAME2 targets were returned
557         assertThat(targets).containsExactly(target1, target2);
558     }
559 
560     @Test
561     public void testGetInstrumentationTargets_runnerFilterSingleFilterSingleResult()
562             throws DeviceNotAvailableException {
563         // Prepare some test data
564         InstrumentationTarget target1 = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
565         InstrumentationTarget target2 = new InstrumentationTarget(PACKAGE_NAME2, RUNNER_NAME2, "");
566         InstrumentationTarget target3 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME1, "");
567         InstrumentationTarget target4 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME3, "");
568 
569         // Set up mocks
570         doReturn(ImmutableList.of(target1, target2, target3, target4))
571                 .when(mInstrumentationParser)
572                 .getInstrumentationTargets();
573         mCoverageTest.setRunnerFilter(ImmutableList.of(RUNNER_NAME2));
574 
575         // Get the instrumentation targets
576         Collection<InstrumentationTarget> targets = mCoverageTest.getInstrumentationTargets();
577 
578         // Verify that only the RUNNER_NAME2 target was returned
579         assertThat(targets).containsExactly(target2);
580     }
581 
582     @Test
583     public void testGetInstrumentationTargets_runnerFilterSingleFilterMultipleResults()
584             throws DeviceNotAvailableException {
585         // Prepare some test data
586         InstrumentationTarget target1 = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
587         InstrumentationTarget target2 = new InstrumentationTarget(PACKAGE_NAME2, RUNNER_NAME2, "");
588         InstrumentationTarget target3 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME1, "");
589         InstrumentationTarget target4 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME3, "");
590 
591         // Set up mocks
592         doReturn(ImmutableList.of(target1, target2, target3, target4))
593                 .when(mInstrumentationParser)
594                 .getInstrumentationTargets();
595         mCoverageTest.setRunnerFilter(ImmutableList.of(RUNNER_NAME1));
596 
597         // Get the instrumentation targets
598         Collection<InstrumentationTarget> targets = mCoverageTest.getInstrumentationTargets();
599 
600         // Verify that both RUNNER_NAME1 targets were returned
601         assertThat(targets).containsExactly(target1, target3);
602     }
603 
604     @Test
605     public void testGetInstrumentationTargets_runnerFilterMultipleFilters()
606             throws DeviceNotAvailableException {
607         // Prepare some test data
608         InstrumentationTarget target1 = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
609         InstrumentationTarget target2 = new InstrumentationTarget(PACKAGE_NAME2, RUNNER_NAME2, "");
610         InstrumentationTarget target3 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME1, "");
611         InstrumentationTarget target4 = new InstrumentationTarget(PACKAGE_NAME3, RUNNER_NAME3, "");
612 
613         // Set up mocks
614         doReturn(ImmutableList.of(target1, target2, target3, target4))
615                 .when(mInstrumentationParser)
616                 .getInstrumentationTargets();
617         mCoverageTest.setRunnerFilter(ImmutableList.of(RUNNER_NAME2, RUNNER_NAME3));
618 
619         // Get the instrumentation targets
620         Collection<InstrumentationTarget> targets = mCoverageTest.getInstrumentationTargets();
621 
622         // Verify that the RUNNER_NAME2 and RUNNER_NAME3 targets were returned
623         assertThat(targets).containsExactly(target2, target4);
624     }
625 
626     @Test
627     public void testDoesRunnerSupportSharding_true() throws DeviceNotAvailableException {
628         // Prepare some test data
629         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
630 
631         // Set up mocks. Return fewer tests when sharding is enabled.
632         doReturn(ImmutableList.of(target)).when(mInstrumentationParser).getInstrumentationTargets();
633         doReturn(FOO_TESTS).when(mTests).get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(0), eq(1));
634         doReturn(Lists.partition(FOO_TESTS, 2).get(0))
635                 .when(mTests)
636                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(0), gt(1));
637 
638         // Verify that the method returns true
639         assertThat(mCoverageTest.doesRunnerSupportSharding(target)).isTrue();
640     }
641 
642     @Test
643     public void testDoesRunnerSupportSharding_false() throws DeviceNotAvailableException {
644         // Prepare some test data
645         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
646 
647         // Set up mocks. Return the same number of tests for any number of shards.
648         doReturn(ImmutableList.of(target)).when(mInstrumentationParser).getInstrumentationTargets();
649         doReturn(FOO_TESTS)
650                 .when(mTests)
651                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
652 
653         // Verify that the method returns false
654         assertThat(mCoverageTest.doesRunnerSupportSharding(target)).isFalse();
655     }
656 
657     @Test
658     public void testGetNumberOfShards() throws DeviceNotAvailableException {
659         // Prepare some test data
660         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
661         List<TestDescription> tests = new ArrayList<>();
662         for (int i = 0; i < 10; i++) {
663             tests.add(new TestDescription(PACKAGE_NAME1, String.format("test%d", i)));
664         }
665 
666         // Set up mocks
667         doReturn(tests).when(mTests).get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
668         mCoverageTest.setMaxTestsPerChunk(1);
669 
670         // Verify that each test will run in a separate shard
671         assertThat(mCoverageTest.getNumberOfShards(target)).isEqualTo(tests.size());
672     }
673 
674     @Test
675     public void testGetNumberOfShards_allTestsInSingleShard() throws DeviceNotAvailableException {
676         // Prepare some test data
677         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
678         List<TestDescription> tests = new ArrayList<>();
679         for (int i = 0; i < 10; i++) {
680             tests.add(new TestDescription(PACKAGE_NAME1, String.format("test%d", i)));
681         }
682 
683         // Set up mocks
684         doReturn(tests).when(mTests).get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
685         mCoverageTest.setMaxTestsPerChunk(10);
686 
687         // Verify that all tests will run in a single shard
688         assertThat(mCoverageTest.getNumberOfShards(target)).isEqualTo(1);
689     }
690 
691     @Test
692     public void testCollectTests() throws DeviceNotAvailableException {
693         // Prepare some test data
694         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
695         List<TestDescription> tests = new ArrayList<>();
696         for (int i = 0; i < 10; i++) {
697             tests.add(new TestDescription(PACKAGE_NAME1, String.format("test%d", i)));
698         }
699 
700         // Set up mocks
701         doReturn(tests).when(mTests).get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), anyInt(), anyInt());
702 
703         // Collect the tests
704         Collection<TestDescription> collectedTests = mCoverageTest.collectTests(target, 0, 1);
705 
706         // Verify that all of the tests were returned
707         assertThat(collectedTests).containsExactlyElementsIn(tests);
708     }
709 
710     @Test
711     public void testCollectTests_withShards() throws DeviceNotAvailableException {
712         // Prepare some test data
713         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
714         int numShards = 3;
715         List<TestDescription> tests = new ArrayList<>();
716         for (int i = 0; i < 10; i++) {
717             tests.add(new TestDescription(PACKAGE_NAME1, String.format("test%d", i)));
718         }
719 
720         // Set up mocks
721         mCoverageTest.setMaxTestsPerChunk((int) Math.ceil((double) tests.size() / numShards));
722         doReturn(tests).when(mTests).get(PACKAGE_NAME1, RUNNER_NAME1, 0, 1);
723         doReturn(tests.subList(0, tests.size() / 2))
724                 .when(mTests)
725                 .get(eq(PACKAGE_NAME1), eq(RUNNER_NAME1), eq(0), eq(2));
726 
727         List<List<TestDescription>> shards =
728                 Lists.partition(tests, (int) Math.ceil((double) tests.size() / numShards));
729         int currentIndex = 0;
730         for (List<TestDescription> shard : shards) {
731             doReturn(shard)
732                     .when(mTests)
733                     .get(PACKAGE_NAME1, RUNNER_NAME1, currentIndex, shards.size());
734             currentIndex++;
735         }
736 
737         // Collect the tests in shards
738         ArrayList<TestDescription> allCollectedTests = new ArrayList<>();
739         for (int shardIndex = 0; shardIndex < numShards; shardIndex++) {
740             // Verify that each shard contains some tests
741             Collection<TestDescription> collectedTests =
742                     mCoverageTest.collectTests(target, shardIndex, numShards);
743             assertThat(collectedTests).containsExactlyElementsIn(shards.get(shardIndex));
744 
745             allCollectedTests.addAll(collectedTests);
746         }
747         // Verify that all of the tests were returned in the end
748         assertThat(allCollectedTests).containsExactlyElementsIn(tests);
749     }
750 
751     @Test
752     public void testCreateTestRunner() {
753         // Prepare some test data
754         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
755 
756         // Create a test runner
757         IRemoteAndroidTestRunner runner = mCoverageTest.createTestRunner(target, 0, 1);
758 
759         // Verify that the runner has the correct values
760         assertThat(runner.getPackageName()).isEqualTo(PACKAGE_NAME1);
761         assertThat(runner.getRunnerName()).isEqualTo(RUNNER_NAME1);
762     }
763 
764     @Test
765     public void testCreateTestRunner_withArgs() {
766         // Prepare some test data
767         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
768         Map<String, String> args = ImmutableMap.of("arg1", "value1", "arg2", "value2");
769 
770         // Set up mocks
771         mCoverageTest.setInstrumentationArgs(args);
772 
773         // Create a test runner
774         FakeTestRunner runner = (FakeTestRunner) mCoverageTest.createTestRunner(target, 0, 1);
775 
776         // Verify that the addInstrumentationArg(..) method was called with each argument
777         assertThat(runner.getArgs()).containsExactlyEntriesIn(args);
778     }
779 
780     @Test
781     public void testCreateTestRunner_withShards() {
782         // Prepare some test data
783         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
784         int shardIndex = 3;
785         int numShards = 5;
786 
787         // Create a test runner
788         FakeTestRunner runner =
789                 (FakeTestRunner) mCoverageTest.createTestRunner(target, shardIndex, numShards);
790 
791         // Verify that the addInstrumentationArg(..) method was called to configure the shards
792         assertThat(runner.getArgs()).containsEntry("shardIndex", Integer.toString(shardIndex));
793         assertThat(runner.getArgs()).containsEntry("numShards", Integer.toString(numShards));
794     }
795 
796     @Test
797     public void testCreateCoverageTest() {
798         // Prepare some test data
799         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
800 
801         // Create a CodeCoverageTest instance
802         InstrumentationTest test = mCoverageTest.createTest(target);
803 
804         // Verify that the test has the correct values
805         assertThat(test.getPackageName()).isEqualTo(PACKAGE_NAME1);
806         assertThat(test.getRunnerName()).isEqualTo(RUNNER_NAME1);
807         assertThat(test.getInstrumentationArg("coverage")).isEqualTo("true");
808     }
809 
810     @Test
811     public void testCreateCoverageTest_withArgs() {
812         // Prepare some test data
813         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
814         Map<String, String> args = ImmutableMap.of("arg1", "value1", "arg2", "value2");
815 
816         // Set up mocks
817         mCoverageTest.setInstrumentationArgs(args);
818 
819         // Create a CodeCoverageTest instance
820         InstrumentationTest test = mCoverageTest.createTest(target, 0, 3);
821 
822         // Verify that the test has the correct values
823         for (Map.Entry<String, String> entry : args.entrySet()) {
824             assertThat(test.getInstrumentationArg(entry.getKey())).isEqualTo(entry.getValue());
825         }
826     }
827 
828     @Test
829     public void testCreateCoverageTest_withShards() {
830         // Prepare some test data
831         InstrumentationTarget target = new InstrumentationTarget(PACKAGE_NAME1, RUNNER_NAME1, "");
832         int shardIndex = 3;
833         int numShards = 5;
834 
835         // Create a CodeCoverageTest instance
836         InstrumentationTest test = mCoverageTest.createTest(target, shardIndex, numShards);
837 
838         // Verify that the addInstrumentationArg(..) method was called to configure the shards
839         assertThat(test.getInstrumentationArg("shardIndex"))
840                 .isEqualTo(Integer.toString(shardIndex));
841         assertThat(test.getInstrumentationArg("numShards")).isEqualTo(Integer.toString(numShards));
842         assertThat(test.getInstrumentationArg("coverage")).isEqualTo("true");
843     }
844 
845     private static final Answer<Void> CALL_RUNNER =
846             invocation -> {
847                 Object[] args = invocation.getArguments();
848                 ((IRemoteAndroidTestRunner) args[0])
849                         .run(
850                                 new TestRunToTestInvocationForwarder(
851                                         (ITestLifeCycleReceiver) args[1]));
852                 return null;
853             };
854 
855     /**
856      * A fake {@link IRemoteAndroidTestRunner} which simulates a test run by notifying the {@link
857      * ITestRunListener}s but does not actually run anything.
858      */
859     private class FakeTestRunner extends RemoteAndroidTestRunner {
860         private Map<String, String> mArgs = new HashMap<>();
861 
862         FakeTestRunner(String packageName, String runnerName) {
863             super(packageName, runnerName, null);
864         }
865 
866         @Override
867         public void addInstrumentationArg(String name, String value) {
868             super.addInstrumentationArg(name, value);
869             mArgs.put(name, value);
870         }
871 
872         Map<String, String> getArgs() {
873             return mArgs;
874         }
875 
876         @Override
877         public void run(Collection<ITestRunListener> listeners) {
878             int shardIndex = Integer.parseInt(getArgs().getOrDefault("shardIndex", "0"));
879             int numShards = Integer.parseInt(getArgs().getOrDefault("numShards", "1"));
880             List<TestDescription> tests =
881                     mTests.get(getPackageName(), getRunnerName(), shardIndex, numShards);
882 
883             // Start the test run
884             listeners.stream().forEach(l -> l.testRunStarted(getPackageName(), tests.size()));
885 
886             // Run each of the tests
887             for (TestDescription test : tests) {
888                 listeners
889                         .stream()
890                         .forEach(l -> l.testStarted(TestDescription.convertToIdentifier(test)));
891                 listeners
892                         .stream()
893                         .forEach(
894                                 l ->
895                                         l.testEnded(
896                                                 TestDescription.convertToIdentifier(test),
897                                                 ImmutableMap.of()));
898             }
899 
900             // Mock out the coverage measurement if necessary
901             Map<String, String> metrics = new HashMap<>();
902             if (getArgs().getOrDefault("coverage", "false").equals("true")) {
903                 String devicePath = String.format(COVERAGE_PATH, getPackageName());
904                 ByteString measurement =
905                         mMeasurements.get(getPackageName(), getRunnerName(), shardIndex, numShards);
906                 mockDeviceFile(devicePath, measurement);
907                 metrics.put(CodeCoverageTest.COVERAGE_REMOTE_FILE_LABEL, devicePath);
908             }
909 
910             // End the test run
911             listeners.stream().forEach(l -> l.testRunEnded(TEST_RUN_TIME, metrics));
912         }
913     }
914 
915     private void mockDeviceFile(String devicePath, ByteString contents) {
916         Answer<File> pullFile =
917                 unused -> {
918                     File ret = mFolder.newFile();
919                     contents.writeTo(new FileOutputStream(ret));
920                     return ret;
921                 };
922         try {
923             doReturn(true).when(mDevice).doesFileExist(devicePath);
924             doAnswer(pullFile).when(mDevice).pullFile(devicePath);
925         } catch (DeviceNotAvailableException impossible) {
926             // Mocks won't actually throw.
927             throw new AssertionError(impossible);
928         }
929     }
930 }
931