From b1a5514200e1f86c3b29262fb7f7635385008bc6 Mon Sep 17 00:00:00 2001
From: Aleksei Chernykh <Aleksei.Chernykh@williamhill.com.au>
Date: Fri, 12 Jan 2018 14:07:34 +1100
Subject: [PATCH] Running tests from a shard in a single instrumentation

---
 .../com/squareup/spoon/SpoonDeviceRunner.java | 52 +++++++++----------
 1 file changed, 25 insertions(+), 27 deletions(-)

diff --git a/spoon-runner/src/main/java/com/squareup/spoon/SpoonDeviceRunner.java b/spoon-runner/src/main/java/com/squareup/spoon/SpoonDeviceRunner.java
index fb88b302..8f36b97c 100644
--- a/spoon-runner/src/main/java/com/squareup/spoon/SpoonDeviceRunner.java
+++ b/spoon-runner/src/main/java/com/squareup/spoon/SpoonDeviceRunner.java
@@ -1,6 +1,6 @@
 package com.squareup.spoon;
 
-import com.android.annotations.Nullable;
+import com.android.annotations.NonNull;
 import com.android.ddmlib.AndroidDebugBridge;
 import com.android.ddmlib.CollectingOutputReceiver;
 import com.android.ddmlib.DdmPreferences;
@@ -298,47 +298,31 @@ private LogRecordingTestRunListener queryTestSet(final String testPackage,
     logDebug(debug, "Querying a list of tests on [%s]", serial);
     RemoteAndroidTestRunner runner = createConfiguredRunner(testPackage, testRunner, device);
     runner.addBooleanArg("log", true);
-    // Add the sharding instrumentation arguments if necessary
-    if (numShards != 0) {
-      addShardingInstrumentationArgs(runner);
-    }
-    if (!isNullOrEmpty(className)) {
-      if (isNullOrEmpty(methodName)) {
-        runner.setClassName(className);
-      } else {
-        runner.setMethodName(className, methodName);
-      }
-    }
-    if (testSize != null) {
-      runner.setTestSize(testSize);
-    }
     runner.run(recorder);
     return recorder;
   }
 
   private void runAllTestOnDevice(final String testPackage, final String testRunner,
       final IDevice device, final List<ITestRunListener> listeners) throws Exception {
-
-    runTestOnDevice(testPackage, testRunner, device, listeners, null);
+    logDebug(debug, "Running tests [%s]", serial);
+    RemoteAndroidTestRunner runner = createConfiguredRunner(testPackage, testRunner, device);
+    runner.removeInstrumentationArg("package");
+    if (codeCoverage) {
+      addCodeCoverageInstrumentationArgs(runner, device);
+    }
+    runner.run(listeners);
   }
 
   private void runTestOnDevice(final String testPackage, final String testRunner,
       final IDevice device, final List<ITestRunListener> listeners,
-      @Nullable final TestIdentifier test) throws Exception {
-
-    if (test != null) {
-      logDebug(debug, "Running %s [%s]", test, serial);
-    } else {
-      logDebug(debug, "Running tests [%s]", serial);
-    }
+      @NonNull final TestIdentifier test) throws Exception {
+    logDebug(debug, "Running %s [%s]", test, serial);
     RemoteAndroidTestRunner runner = createConfiguredRunner(testPackage, testRunner, device);
     runner.removeInstrumentationArg("package");
     if (codeCoverage) {
       addCodeCoverageInstrumentationArgs(runner, device);
     }
-    if (test != null) {
-      runner.setMethodName(test.getClassName(), test.getTestName());
-    }
+    runner.setMethodName(test.getClassName(), test.getTestName());
     runner.run(listeners);
   }
 
@@ -356,6 +340,20 @@ private RemoteAndroidTestRunner createConfiguredRunner(String testPackage, Strin
       }
       runner.addInstrumentationArg(entry.getKey(), entry.getValue());
     }
+    // Add the sharding instrumentation arguments if necessary
+    if (numShards != 0) {
+      addShardingInstrumentationArgs(runner);
+    }
+    if (!isNullOrEmpty(className)) {
+      if (isNullOrEmpty(methodName)) {
+        runner.setClassName(className);
+      } else {
+        runner.setMethodName(className, methodName);
+      }
+    }
+    if (testSize != null) {
+      runner.setTestSize(testSize);
+    }
     return runner;
   }