diff --git a/tools/tradefed-host/src/com/android/cts/tradefed/testtype/CtsTest.java b/tools/tradefed-host/src/com/android/cts/tradefed/testtype/CtsTest.java
index 2834793..f49c4a0 100644
--- a/tools/tradefed-host/src/com/android/cts/tradefed/testtype/CtsTest.java
+++ b/tools/tradefed-host/src/com/android/cts/tradefed/testtype/CtsTest.java
@@ -32,7 +32,6 @@
 import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.testtype.IResumableTest;
 import com.android.tradefed.testtype.IShardableTest;
-import com.android.tradefed.testtype.InstrumentationTest;
 import com.android.tradefed.util.xml.AbstractXmlParser.ParseException;
 
 import java.io.BufferedInputStream;
@@ -119,6 +118,7 @@
     private List<KnownTests> mRemainingTests = null;
 
     private CtsBuildHelper mCtsBuild = null;
+    private IBuildInfo mBuildInfo = null;
 
     /**
      * {@inheritDoc}
@@ -207,6 +207,7 @@
                     build.getClass().getName()));
         }
         try {
+            mBuildInfo = build;
             mCtsBuild = new CtsBuildHelper((IFolderBuildInfo)build);
             mCtsBuild.validateStructure();
         } catch (FileNotFoundException e) {
@@ -249,12 +250,10 @@
             if (test instanceof IDeviceTest) {
                 ((IDeviceTest)test).setDevice(getDevice());
             }
-            // Increment the timeout for collecting the tests.
-            // TODO: move this to testPackage.createTest() instead and only increase timeout when
-            // tests number is large.
-            if (test instanceof InstrumentationTest) {
-                ((InstrumentationTest)test).setCollectsTestsShellTimeout(10*60*1000);
+            if (test instanceof IBuildReceiver) {
+                ((IBuildReceiver)test).setBuild(mBuildInfo);
             }
+
             ResultFilter filter = new ResultFilter(listener, testPair.getKnownTests());
             test.run(filter);
             mRemainingTests.remove(0);
diff --git a/tools/tradefed-host/src/com/android/cts/tradefed/testtype/ITestPackageDef.java b/tools/tradefed-host/src/com/android/cts/tradefed/testtype/ITestPackageDef.java
index e4f13b5..f4258d1 100644
--- a/tools/tradefed-host/src/com/android/cts/tradefed/testtype/ITestPackageDef.java
+++ b/tools/tradefed-host/src/com/android/cts/tradefed/testtype/ITestPackageDef.java
@@ -27,7 +27,7 @@
  * <p/>
  * Knows how to translate this info into a runnable {@link IRemoteTest}.
  */
-interface ITestPackageDef {
+public interface ITestPackageDef {
 
     /**
      * Get the unique URI of the test package.
diff --git a/tools/tradefed-host/src/com/android/cts/tradefed/testtype/InstrumentationApkTest.java b/tools/tradefed-host/src/com/android/cts/tradefed/testtype/InstrumentationApkTest.java
new file mode 100644
index 0000000..b3d2eec
--- /dev/null
+++ b/tools/tradefed-host/src/com/android/cts/tradefed/testtype/InstrumentationApkTest.java
@@ -0,0 +1,104 @@
+/*
+ * Copyright (C) 2011 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.cts.tradefed.testtype;
+
+import com.android.cts.tradefed.build.CtsBuildHelper;
+import com.android.ddmlib.Log;
+import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.build.IFolderBuildInfo;
+import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.result.ITestInvocationListener;
+import com.android.tradefed.testtype.IBuildReceiver;
+import com.android.tradefed.testtype.InstrumentationTest;
+
+import java.io.FileNotFoundException;
+import java.util.ArrayList;
+import java.util.Collection;
+
+import junit.framework.Assert;
+
+/**
+ * A {@link InstrumentationTest] that will install CTS apks before test execution,
+ * and uninstall on execution completion.
+ */
+public class InstrumentationApkTest extends InstrumentationTest implements IBuildReceiver {
+
+    private static final String LOG_TAG = "InstrumentationApkTest";
+
+    /** the file names of the CTS apks to install */
+    private Collection<String> mInstallFileNames = new ArrayList<String>();
+    private Collection<String> mUninstallPackages = new ArrayList<String>();
+
+    private CtsBuildHelper mCtsBuild = null;
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void setBuild(IBuildInfo build) {
+        if (!(build instanceof IFolderBuildInfo)) {
+            throw new IllegalArgumentException(String.format(
+                    "Wrong build type. Expected %s, received %s", IFolderBuildInfo.class.getName(),
+                    build.getClass().getName()));
+        }
+        try {
+            mCtsBuild  = new CtsBuildHelper((IFolderBuildInfo)build);
+            mCtsBuild.validateStructure();
+        } catch (FileNotFoundException e) {
+            throw new IllegalArgumentException("Invalid CTS build provided.", e);
+        }
+    }
+
+    /**
+     * Add an apk to install.
+     *
+     * @param apkFileName the apk file name
+     * @param packageName the apk's Android package name
+     */
+    public void addInstallApk(String apkFileName, String packageName) {
+        mInstallFileNames.add(apkFileName);
+        mUninstallPackages.add(packageName);
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void run(final ITestInvocationListener listener)
+            throws DeviceNotAvailableException {
+        Assert.assertNotNull("missing device", getDevice());
+        Assert.assertNotNull("missing build", mCtsBuild);
+
+        try {
+            for (String apkFileName : mInstallFileNames) {
+                Log.d(LOG_TAG, String.format("Installing %s on %s", apkFileName,
+                        getDevice().getSerialNumber()));
+                try {
+                    getDevice().installPackage(mCtsBuild.getTestApp(apkFileName), true);
+                } catch (FileNotFoundException e) {
+                    Log.e(LOG_TAG, e);
+                }
+            }
+            super.run(listener);
+        } finally {
+            for (String packageName : mUninstallPackages) {
+                Log.d(LOG_TAG, String.format("Uninstalling %s on %s", packageName,
+                        getDevice().getSerialNumber()));
+                getDevice().uninstallPackage(packageName);
+            }
+        }
+    }
+}
diff --git a/tools/tradefed-host/src/com/android/cts/tradefed/testtype/InstrumentationAppTest.java b/tools/tradefed-host/src/com/android/cts/tradefed/testtype/InstrumentationAppTest.java
deleted file mode 100644
index 47fc14a..0000000
--- a/tools/tradefed-host/src/com/android/cts/tradefed/testtype/InstrumentationAppTest.java
+++ /dev/null
@@ -1,72 +0,0 @@
-/*
- * Copyright (C) 2011 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.cts.tradefed.testtype;
-
-import com.android.ddmlib.Log;
-import com.android.tradefed.device.DeviceNotAvailableException;
-import com.android.tradefed.result.ITestInvocationListener;
-import com.android.tradefed.testtype.InstrumentationTest;
-
-import java.io.File;
-import java.util.ArrayList;
-import java.util.Collection;
-
-/**
- * A {@link InstrumentationTest] that will install other dependent apks before test execution,
- * and uninstall apps on execution completion.
- */
-public class InstrumentationAppTest extends InstrumentationTest {
-
-    private static final String LOG_TAG = "InstrumentationAppTest";
-
-    // TODO: consider moving this class to tradefed proper
-
-    private Collection<File> mInstallFiles = new ArrayList<File>();
-    private Collection<String> mInstallPackages = new ArrayList<String>();
-
-    /**
-     * Add a dependent apk to install.
-     *
-     * @param apkFile the apk file
-     * @param packageName the apk's Android package name
-     */
-    public void addInstallApp(File apkFile, String packageName) {
-        mInstallFiles.add(apkFile);
-        mInstallPackages.add(packageName);
-    }
-
-    @Override
-    public void run(final ITestInvocationListener listener)
-            throws DeviceNotAvailableException {
-        if (getDevice() == null) {
-            throw new IllegalStateException("missing device");
-        }
-        try {
-            for (File apkFile : mInstallFiles) {
-                Log.d(LOG_TAG, String.format("Installing %s on %s", apkFile.getName(),
-                        getDevice().getSerialNumber()));
-                getDevice().installPackage(apkFile, true);
-            }
-            super.run(listener);
-        } finally {
-            for (String packageName : mInstallPackages) {
-                Log.d(LOG_TAG, String.format("Uninstalling %s on %s", packageName,
-                        getDevice().getSerialNumber()));
-                getDevice().uninstallPackage(packageName);
-            }
-        }
-    }
-}
diff --git a/tools/tradefed-host/src/com/android/cts/tradefed/testtype/TestPackageDef.java b/tools/tradefed-host/src/com/android/cts/tradefed/testtype/TestPackageDef.java
index fd896f2..6b165be 100644
--- a/tools/tradefed-host/src/com/android/cts/tradefed/testtype/TestPackageDef.java
+++ b/tools/tradefed-host/src/com/android/cts/tradefed/testtype/TestPackageDef.java
@@ -163,24 +163,18 @@
             return instrTest;
         } else if (mIsReferenceAppTest) {
             // a reference app test is just a InstrumentationTest with one extra apk to install
-            InstrumentationAppTest instrTest = new InstrumentationAppTest();
-            File apkFile = new File(testCaseDir, String.format("%s.apk", mApkToTestName));
-            if (!apkFile.exists()) {
-                Log.w(LOG_TAG, String.format("Could not find apk file %s",
-                        apkFile.getAbsolutePath()));
-                return null;
-            }
-            instrTest.addInstallApp(apkFile, mPackageToTest);
-            return setInstrumentationTest(testCaseDir, className, methodName, instrTest);
+            InstrumentationApkTest instrTest = new InstrumentationApkTest();
+            instrTest.addInstallApk(String.format("%s.apk", mApkToTestName), mPackageToTest);
+            return setInstrumentationTest(className, methodName, instrTest);
         } else {
             Log.d(LOG_TAG, String.format("Creating instrumentation test for %s", mName));
-            InstrumentationTest instrTest = new InstrumentationTest();
-            return setInstrumentationTest(testCaseDir, className, methodName, instrTest);
+            InstrumentationApkTest instrTest = new InstrumentationApkTest();
+            return setInstrumentationTest(className, methodName, instrTest);
         }
     }
 
     /**
-     * Populates given {@link InstrumentationTest} with data from the package xml
+     * Populates given {@link InstrumentationApkTest} with data from the package xml
      *
      * @param testCaseDir
      * @param className
@@ -188,20 +182,18 @@
      * @param instrTest
      * @return the populated {@link InstrumentationTest} or <code>null</code>
      */
-    private InstrumentationTest setInstrumentationTest(File testCaseDir, String className,
-            String methodName, InstrumentationTest instrTest) {
+    private InstrumentationApkTest setInstrumentationTest(String className,
+            String methodName, InstrumentationApkTest instrTest) {
         instrTest.setPackageName(mAppNameSpace);
         instrTest.setRunnerName(mRunner);
         instrTest.setClassName(className);
         instrTest.setMethodName(methodName);
         // mName means 'apk file name' for instrumentation tests
-        File apkFile = new File(testCaseDir, String.format("%s.apk", mName));
-        if (!apkFile.exists()) {
-            Log.w(LOG_TAG, String.format("Could not find apk file %s",
-                    apkFile.getAbsolutePath()));
-            return null;
+        instrTest.addInstallApk(String.format("%s.apk", mName), mAppNameSpace);
+        if (mTests.size() > 1000) {
+            // TODO: hack, large test suites can take longer to collect tests, increase timeout
+            instrTest.setCollectsTestsShellTimeout(10*60*1000);
         }
-        instrTest.setInstallFile(apkFile);
         return instrTest;
     }