Ensure cleanUp of build info is called

Ensure we call clean up on the build info even in case of
- exceptions
- multi device where only some of the builds are downloaded

Test: unit tests
Bug: 64576937
Change-Id: I190fab6b7e171acafdf28502df7d04146006441b

4 files changed, 359 insertions, 5 deletions

diff --git a/src/com/android/tradefed/invoker/TestInvocation.java b/src/com/android/tradefed/invoker/TestInvocation.java
index 605f7018b..237a5e9b1 100644
--- a/src/com/android/tradefed/invoker/TestInvocation.java
+++ b/src/com/android/tradefed/invoker/TestInvocation.java
@@ -901,8 +901,6 @@ public class TestInvocation implements ITestInvocation {
                             "No build found to test for device: %s",
                             device.getSerialNumber());
                     rescheduleTest(config, rescheduler);
-                    // save current log contents to global log
-                    getLogRegistry().dumpToGlobalLog(config.getLogOutput());
                     // Set the exit code to error
                     setExitCode(ExitCode.NO_BUILD,
                             new BuildRetrievalError("No build found to test."));
@@ -941,6 +939,22 @@ public class TestInvocation implements ITestInvocation {
         } catch (IOException e) {
             CLog.e(e);
         } finally {
+
+            // Ensure build infos are always cleaned up at the end of invocation.
+            for (String cleanUpDevice : context.getDeviceConfigNames()) {
+                if (context.getBuildInfo(cleanUpDevice) != null) {
+                    try {
+                        config.getDeviceConfigByName(cleanUpDevice)
+                                .getBuildProvider()
+                                .cleanUp(context.getBuildInfo(cleanUpDevice));
+                    } catch (RuntimeException e) {
+                        // We catch an simply log exception in cleanUp to avoid missing any final
+                        // step of the invocation.
+                        CLog.e(e);
+                    }
+                }
+            }
+
             // ensure we always deregister the logger
             for (String deviceName : context.getDeviceConfigNames()) {
                 if (!(context.getDevice(deviceName).getIDevice() instanceof StubDevice)) {
diff --git a/tests/src/com/android/tradefed/UnitTests.java b/tests/src/com/android/tradefed/UnitTests.java
index 759b73b41..bfc22e98a 100644
--- a/tests/src/com/android/tradefed/UnitTests.java
+++ b/tests/src/com/android/tradefed/UnitTests.java
@@ -63,6 +63,7 @@ import com.android.tradefed.device.TopHelperTest;
 import com.android.tradefed.device.WaitDeviceRecoveryTest;
 import com.android.tradefed.device.WifiHelperTest;
 import com.android.tradefed.invoker.InvocationContextTest;
+import com.android.tradefed.invoker.TestInvocationMultiTest;
 import com.android.tradefed.invoker.TestInvocationTest;
 import com.android.tradefed.invoker.shard.ShardHelperTest;
 import com.android.tradefed.invoker.shard.StrictShardHelperTest;
@@ -290,6 +291,7 @@ import org.junit.runners.Suite.SuiteClasses;
 
     // invoker
     InvocationContextTest.class,
+    TestInvocationMultiTest.class,
     TestInvocationTest.class,
 
     // invoker.shard
diff --git a/tests/src/com/android/tradefed/invoker/TestInvocationMultiTest.java b/tests/src/com/android/tradefed/invoker/TestInvocationMultiTest.java
new file mode 100644
index 000000000..2307f1dec
--- /dev/null
+++ b/tests/src/com/android/tradefed/invoker/TestInvocationMultiTest.java
@@ -0,0 +1,329 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.tradefed.invoker;
+
+import com.android.tradefed.build.BuildInfo;
+import com.android.tradefed.build.BuildRetrievalError;
+import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.build.IBuildProvider;
+import com.android.tradefed.command.CommandOptions;
+import com.android.tradefed.command.CommandRunner.ExitCode;
+import com.android.tradefed.config.DeviceConfigurationHolder;
+import com.android.tradefed.config.IConfiguration;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.StubDevice;
+import com.android.tradefed.invoker.shard.IShardHelper;
+import com.android.tradefed.invoker.shard.ShardHelper;
+import com.android.tradefed.log.ILeveledLogOutput;
+import com.android.tradefed.log.ILogRegistry;
+import com.android.tradefed.result.ByteArrayInputStreamSource;
+import com.android.tradefed.result.ILogSaver;
+import com.android.tradefed.result.ITestInvocationListener;
+import com.android.tradefed.result.LogFile;
+
+import org.easymock.EasyMock;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/** Unit tests for {@link TestInvocation} for multi device invocation. */
+@RunWith(JUnit4.class)
+public class TestInvocationMultiTest {
+    private TestInvocation mInvocation;
+    private IInvocationContext mContext;
+    private IConfiguration mMockConfig;
+    private IRescheduler mMockRescheduler;
+    private ITestInvocationListener mMockTestListener;
+    private ILogSaver mMockLogSaver;
+    private ILeveledLogOutput mMockLogger;
+    private ILogRegistry mMockLogRegistry;
+
+    private ITestDevice mDevice1;
+    private ITestDevice mDevice2;
+    private IBuildProvider mProvider1;
+    private IBuildProvider mProvider2;
+
+    @Before
+    public void setUp() {
+        mContext = new InvocationContext();
+        mMockConfig = EasyMock.createMock(IConfiguration.class);
+        mMockRescheduler = EasyMock.createMock(IRescheduler.class);
+        mMockTestListener = EasyMock.createMock(ITestInvocationListener.class);
+        mMockLogSaver = EasyMock.createMock(ILogSaver.class);
+        mMockLogger = EasyMock.createMock(ILeveledLogOutput.class);
+        mMockLogRegistry = EasyMock.createMock(ILogRegistry.class);
+        mInvocation =
+                new TestInvocation() {
+                    @Override
+                    ILogRegistry getLogRegistry() {
+                        return mMockLogRegistry;
+                    }
+
+                    @Override
+                    protected IShardHelper createShardHelper() {
+                        return new ShardHelper();
+                    }
+
+                    @Override
+                    protected void setExitCode(ExitCode code, Throwable stack) {
+                        // empty on purpose
+                    }
+                };
+    }
+
+    private void makeTwoDeviceContext() throws Exception {
+        mDevice1 = EasyMock.createMock(ITestDevice.class);
+        EasyMock.expect(mDevice1.getIDevice()).andStubReturn(new StubDevice("serial1"));
+        EasyMock.expect(mDevice1.getSerialNumber()).andStubReturn("serial1");
+        mDevice1.clearLastConnectedWifiNetwork();
+        DeviceConfigurationHolder holder1 = new DeviceConfigurationHolder();
+        mProvider1 = EasyMock.createMock(IBuildProvider.class);
+        holder1.addSpecificConfig(mProvider1);
+        EasyMock.expect(mMockConfig.getDeviceConfigByName("device1")).andStubReturn(holder1);
+        mDevice1.setOptions(EasyMock.anyObject());
+        mDevice1.setRecovery(EasyMock.anyObject());
+
+        mDevice2 = EasyMock.createMock(ITestDevice.class);
+        EasyMock.expect(mDevice2.getIDevice()).andStubReturn(new StubDevice("serial2"));
+        EasyMock.expect(mDevice2.getSerialNumber()).andStubReturn("serial2");
+        mDevice2.clearLastConnectedWifiNetwork();
+        DeviceConfigurationHolder holder2 = new DeviceConfigurationHolder();
+        mProvider2 = EasyMock.createMock(IBuildProvider.class);
+        holder2.addSpecificConfig(mProvider2);
+        EasyMock.expect(mMockConfig.getDeviceConfigByName("device2")).andStubReturn(holder2);
+        mDevice2.setOptions(EasyMock.anyObject());
+
+        mContext.addAllocatedDevice("device1", mDevice1);
+        mContext.addAllocatedDevice("device2", mDevice2);
+    }
+
+    /**
+     * Test for multi device invocation when the first download succeed and second one is missing.
+     * We clean up all the downloaded builds.
+     */
+    @Test
+    public void testRunBuildProvider_oneMiss() throws Throwable {
+        makeTwoDeviceContext();
+
+        List<ITestInvocationListener> configListener = new ArrayList<>();
+        configListener.add(mMockTestListener);
+        EasyMock.expect(mMockConfig.getTestInvocationListeners())
+                .andReturn(configListener)
+                .times(2);
+        EasyMock.expect(mMockConfig.getProfiler()).andReturn(null);
+        EasyMock.expect(mMockConfig.getLogSaver()).andReturn(mMockLogSaver);
+        EasyMock.expect(mMockConfig.getLogOutput()).andReturn(mMockLogger).times(4);
+        mMockLogger.init();
+        mMockLogger.closeLog();
+
+        mMockLogRegistry.registerLogger(mMockLogger);
+        mMockLogRegistry.dumpToGlobalLog(mMockLogger);
+        mMockLogRegistry.unregisterLogger();
+
+        EasyMock.expect(mMockConfig.getCommandLine()).andStubReturn("empty");
+        EasyMock.expect(mMockConfig.getCommandOptions()).andStubReturn(new CommandOptions());
+        EasyMock.expect(mMockConfig.getTests()).andStubReturn(new ArrayList<>());
+        IBuildInfo build1 = new BuildInfo();
+        EasyMock.expect(mProvider1.getBuild()).andReturn(build1);
+        // Second build is not found
+        EasyMock.expect(mProvider2.getBuild()).andReturn(null);
+        // The downloaded build is cleaned
+        mProvider1.cleanUp(build1);
+
+        EasyMock.replay(
+                mMockConfig,
+                mMockRescheduler,
+                mMockTestListener,
+                mMockLogSaver,
+                mMockLogger,
+                mMockLogRegistry,
+                mDevice1,
+                mDevice2,
+                mProvider1,
+                mProvider2);
+        mInvocation.invoke(
+                mContext, mMockConfig, mMockRescheduler, new ITestInvocationListener[] {});
+        EasyMock.verify(
+                mMockConfig,
+                mMockRescheduler,
+                mMockTestListener,
+                mMockLogSaver,
+                mMockLogger,
+                mMockLogRegistry,
+                mDevice1,
+                mDevice2,
+                mProvider1,
+                mProvider2);
+    }
+
+    @Test
+    public void testRunBuildProvider_oneThrow() throws Throwable {
+        makeTwoDeviceContext();
+
+        List<ITestInvocationListener> configListener = new ArrayList<>();
+        configListener.add(mMockTestListener);
+        EasyMock.expect(mMockConfig.getTestInvocationListeners())
+                .andReturn(configListener)
+                .times(2);
+        EasyMock.expect(mMockConfig.getProfiler()).andReturn(null);
+        EasyMock.expect(mMockConfig.getLogSaver()).andReturn(mMockLogSaver);
+        EasyMock.expect(mMockConfig.getLogOutput()).andStubReturn(mMockLogger);
+        mMockLogger.init();
+        EasyMock.expect(mMockLogger.getLog())
+                .andReturn(new ByteArrayInputStreamSource("fake".getBytes()));
+        mMockLogger.closeLog();
+        EasyMock.expectLastCall().times(2);
+
+        mMockLogRegistry.registerLogger(mMockLogger);
+        mMockLogRegistry.dumpToGlobalLog(mMockLogger);
+        mMockLogRegistry.unregisterLogger();
+        EasyMock.expectLastCall().times(2);
+
+        EasyMock.expect(mMockConfig.getCommandLine()).andStubReturn("empty");
+        EasyMock.expect(mMockConfig.getCommandOptions()).andStubReturn(new CommandOptions());
+        EasyMock.expect(mMockConfig.getTests()).andStubReturn(new ArrayList<>());
+
+        mMockTestListener.invocationStarted(mContext);
+        mMockLogSaver.invocationStarted(mContext);
+        mMockTestListener.invocationFailed(EasyMock.anyObject());
+        mMockTestListener.testLog(EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject());
+        EasyMock.expect(
+                        mMockLogSaver.saveLogData(
+                                EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject()))
+                .andReturn(new LogFile("", "", false, true));
+        mMockTestListener.invocationEnded(EasyMock.anyLong());
+        EasyMock.expect(mMockTestListener.getSummary()).andReturn(null);
+        mMockLogSaver.invocationEnded(EasyMock.anyLong());
+
+        IBuildInfo build1 = new BuildInfo();
+        EasyMock.expect(mProvider1.getBuild()).andReturn(build1);
+        // Second build is not found
+        EasyMock.expect(mProvider2.getBuild()).andThrow(new BuildRetrievalError("fail"));
+        // The downloaded build is cleaned
+        mProvider1.cleanUp(build1);
+        // A second build from the BuildRetrievalError is generated but still cleaned.
+        mProvider2.cleanUp(EasyMock.anyObject());
+
+        EasyMock.replay(
+                mMockConfig,
+                mMockRescheduler,
+                mMockTestListener,
+                mMockLogSaver,
+                mMockLogger,
+                mMockLogRegistry,
+                mDevice1,
+                mDevice2,
+                mProvider1,
+                mProvider2);
+        mInvocation.invoke(
+                mContext, mMockConfig, mMockRescheduler, new ITestInvocationListener[] {});
+        EasyMock.verify(
+                mMockConfig,
+                mMockRescheduler,
+                mMockTestListener,
+                mMockLogSaver,
+                mMockLogger,
+                mMockLogRegistry,
+                mDevice1,
+                mDevice2,
+                mProvider1,
+                mProvider2);
+    }
+
+    /**
+     * Test when the provider clean up throws an exception, we still continue to clean up the rest
+     * to ensure nothing is left afterward.
+     */
+    @Test
+    public void testRunBuildProvider_cleanUpThrow() throws Throwable {
+        makeTwoDeviceContext();
+
+        List<ITestInvocationListener> configListener = new ArrayList<>();
+        configListener.add(mMockTestListener);
+        EasyMock.expect(mMockConfig.getTestInvocationListeners())
+                .andReturn(configListener)
+                .times(2);
+        EasyMock.expect(mMockConfig.getProfiler()).andReturn(null);
+        EasyMock.expect(mMockConfig.getLogSaver()).andReturn(mMockLogSaver);
+        EasyMock.expect(mMockConfig.getLogOutput()).andStubReturn(mMockLogger);
+        mMockLogger.init();
+        EasyMock.expect(mMockLogger.getLog())
+                .andReturn(new ByteArrayInputStreamSource("fake".getBytes()));
+        mMockLogger.closeLog();
+        EasyMock.expectLastCall().times(2);
+
+        mMockLogRegistry.registerLogger(mMockLogger);
+        mMockLogRegistry.dumpToGlobalLog(mMockLogger);
+        mMockLogRegistry.unregisterLogger();
+        EasyMock.expectLastCall().times(2);
+
+        EasyMock.expect(mMockConfig.getCommandLine()).andStubReturn("empty");
+        EasyMock.expect(mMockConfig.getCommandOptions()).andStubReturn(new CommandOptions());
+        EasyMock.expect(mMockConfig.getTests()).andStubReturn(new ArrayList<>());
+
+        mMockTestListener.invocationStarted(mContext);
+        mMockLogSaver.invocationStarted(mContext);
+        mMockTestListener.invocationFailed(EasyMock.anyObject());
+        mMockTestListener.testLog(EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject());
+        EasyMock.expect(
+                        mMockLogSaver.saveLogData(
+                                EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject()))
+                .andReturn(new LogFile("", "", false, true));
+        mMockTestListener.invocationEnded(EasyMock.anyLong());
+        EasyMock.expect(mMockTestListener.getSummary()).andReturn(null);
+        mMockLogSaver.invocationEnded(EasyMock.anyLong());
+
+        IBuildInfo build1 = new BuildInfo();
+        EasyMock.expect(mProvider1.getBuild()).andReturn(build1);
+        // Second build is not found
+        EasyMock.expect(mProvider2.getBuild()).andThrow(new BuildRetrievalError("fail"));
+        // The downloaded build is cleaned
+        mProvider1.cleanUp(build1);
+        EasyMock.expectLastCall().andThrow(new RuntimeException("I failed to clean!"));
+        // A second build from the BuildRetrievalError is generated but still cleaned, even if the
+        // first clean up failed.
+        mProvider2.cleanUp(EasyMock.anyObject());
+
+        EasyMock.replay(
+                mMockConfig,
+                mMockRescheduler,
+                mMockTestListener,
+                mMockLogSaver,
+                mMockLogger,
+                mMockLogRegistry,
+                mDevice1,
+                mDevice2,
+                mProvider1,
+                mProvider2);
+        mInvocation.invoke(
+                mContext, mMockConfig, mMockRescheduler, new ITestInvocationListener[] {});
+        EasyMock.verify(
+                mMockConfig,
+                mMockRescheduler,
+                mMockTestListener,
+                mMockLogSaver,
+                mMockLogger,
+                mMockLogRegistry,
+                mDevice1,
+                mDevice2,
+                mProvider1,
+                mProvider2);
+    }
+}
diff --git a/tests/src/com/android/tradefed/invoker/TestInvocationTest.java b/tests/src/com/android/tradefed/invoker/TestInvocationTest.java
index d298d865b..963b9854c 100644
--- a/tests/src/com/android/tradefed/invoker/TestInvocationTest.java
+++ b/tests/src/com/android/tradefed/invoker/TestInvocationTest.java
@@ -313,6 +313,7 @@ public class TestInvocationTest extends TestCase {
         mMockLogRegistry.unregisterLogger();
         mMockLogRegistry.dumpToGlobalLog(mMockLogger);
         mMockLogger.closeLog();
+        mMockBuildProvider.cleanUp(mMockBuildInfo);
         replayMocks(test, mockRescheduler);
         mTestInvocation.invoke(mStubInvocationMetadata, mStubConfiguration, mockRescheduler);
         verifyMocks(test, mockRescheduler);
@@ -327,8 +328,8 @@ public class TestInvocationTest extends TestCase {
         EasyMock.expect(mMockBuildProvider.getBuild()).andReturn(null);
         IRemoteTest test = EasyMock.createMock(IRemoteTest.class);
         mStubConfiguration.setTest(test);
+        mMockBuildProvider.cleanUp(mMockBuildInfo);
         mMockLogRegistry.dumpToGlobalLog(mMockLogger);
-        EasyMock.expectLastCall().times(2);
         setupInvoke();
         replayMocks(test, mockRescheduler);
         mTestInvocation.invoke(mStubInvocationMetadata, mStubConfiguration, mockRescheduler);
@@ -349,8 +350,9 @@ public class TestInvocationTest extends TestCase {
         mStubConfiguration.setTest(test);
         mStubConfiguration.getCommandOptions().setLoopMode(false);
         mMockLogRegistry.dumpToGlobalLog(mMockLogger);
-        EasyMock.expectLastCall().times(2);
+        EasyMock.expectLastCall().times(1);
         setupInvoke();
+        mMockBuildProvider.cleanUp(mMockBuildInfo);
         replayMocks(test);
         EasyMock.replay(mockRescheduler);
         mTestInvocation.invoke(mStubInvocationMetadata, mStubConfiguration, mockRescheduler);
@@ -550,6 +552,7 @@ public class TestInvocationTest extends TestCase {
         EasyMock.expect(mockRescheduler.scheduleConfig(EasyMock.capture(capturedConfig)))
                 .andReturn(Boolean.TRUE);
         mMockBuildProvider.cleanUp(mMockBuildInfo);
+        EasyMock.expectLastCall().times(2);
         mMockDevice.clearLastConnectedWifiNetwork();
         mMockDevice.stopLogcat();
 
@@ -614,6 +617,7 @@ public class TestInvocationTest extends TestCase {
         mMockLogSaver.invocationEnded(EasyMock.anyLong());
         EasyMock.expect(resumeListener.getSummary()).andReturn(null);
         mMockBuildInfo.cleanUp();
+        EasyMock.expectLastCall().times(2);
         mMockLogger.closeLog();
         EasyMock.expectLastCall().times(3);
         mMockDevice.clearLastConnectedWifiNetwork();
@@ -947,6 +951,7 @@ public class TestInvocationTest extends TestCase {
         EasyMock.expect(mMockLogger.getLog()).andReturn(EMPTY_STREAM_SOURCE);
         mMockBuildInfo.setDeviceSerial(SERIAL);
         mMockBuildProvider.cleanUp(mMockBuildInfo);
+        EasyMock.expectLastCall().times(2);
         setupMockSuccessListeners();
         EasyMock.expect(mMockBuildProvider.getBuild()).andReturn(mMockBuildInfo);
         mMockBuildInfo.addBuildAttribute("command_line_args", "run empty");
@@ -977,6 +982,7 @@ public class TestInvocationTest extends TestCase {
         EasyMock.expect(mMockLogger.getLog()).andReturn(EMPTY_STREAM_SOURCE);
         mMockBuildInfo.setDeviceSerial(SERIAL);
         mMockBuildProvider.cleanUp(mMockBuildInfo);
+        EasyMock.expectLastCall().times(2);
         setupMockSuccessListeners();
         EasyMock.expect(mMockBuildProvider.getBuild()).andReturn(mMockBuildInfo);
         mMockBuildInfo.addBuildAttribute("command_line_args", "run empty");
@@ -1004,6 +1010,7 @@ public class TestInvocationTest extends TestCase {
         EasyMock.expect(mMockLogger.getLog()).andReturn(EMPTY_STREAM_SOURCE);
         mMockBuildInfo.setDeviceSerial(SERIAL);
         mMockBuildProvider.cleanUp(mMockBuildInfo);
+        EasyMock.expectLastCall().times(2);
         setupMockSuccessListeners();
         EasyMock.expect(mMockBuildProvider.getBuild()).andReturn(mMockBuildInfo);
         mMockBuildInfo.addBuildAttribute("command_line_args", "run empty");
@@ -1053,6 +1060,7 @@ public class TestInvocationTest extends TestCase {
         mockProvider.setInvocationContext((IInvocationContext)EasyMock.anyObject());
         EasyMock.expect(mockProvider.getBuild(mMockDevice)).andReturn(mMockBuildInfo);
         mockProvider.cleanUp(mMockBuildInfo);
+        EasyMock.expectLastCall().times(2);
         mMockLogRegistry.dumpToGlobalLog(mMockLogger);
         mMockLogRegistry.unregisterLogger();
         mMockLogger.closeLog();
@@ -1122,6 +1130,7 @@ public class TestInvocationTest extends TestCase {
         EasyMock.expect(mMockLogger.getLog()).andReturn(EMPTY_STREAM_SOURCE);
         mMockBuildInfo.setDeviceSerial(SERIAL);
         mMockBuildProvider.cleanUp(mMockBuildInfo);
+        EasyMock.expectLastCall().anyTimes();
         mMockBuildInfo.setTestTag(EasyMock.eq("stub"));
         EasyMock.expectLastCall();
         EasyMock.expect(mMockBuildInfo.getTestTag()).andStubReturn("");
@@ -1290,7 +1299,7 @@ public class TestInvocationTest extends TestCase {
         EasyMock.expect(test.split()).andReturn(shards);
         mStubConfiguration.setTest(test);
         mStubConfiguration.setCommandLine(commandLine);
-
+        mMockBuildProvider.cleanUp(mMockBuildInfo);
         setupInvoke();
         setupNShardInvocation(shardCount, command);
         mMockLogRegistry.dumpToGlobalLog(mMockLogger);