Added ability to get pids by command only.

This previously was using grep, and checking every part of the command
for the identifier, but now only uses the starting command.

Bug: None
Test: None
Change-Id: I98ddefeb84cca64c33975c0bc3f4058255220bbe

4 files changed, 36 insertions, 4 deletions

diff --git a/tools/lab/metrics/process_time_metric.py b/tools/lab/metrics/process_time_metric.py
index db62af1e53..bc7194b342 100644
--- a/tools/lab/metrics/process_time_metric.py
+++ b/tools/lab/metrics/process_time_metric.py
@@ -89,7 +89,7 @@ class ProcessTimeMetric(Metric):
           A list of PID strings
         """
         # Get ids of processes with 'adb' or 'fastboot' in name
-        adb_result = self._shell.get_pids('adb')
-        fastboot_result = self._shell.get_pids('fastboot')
+        adb_result = self._shell.get_command_pids('adb')
+        fastboot_result = self._shell.get_command_pids('fastboot')
         # concatenate two generator objects, return as list
         return list(itertools.chain(adb_result, fastboot_result))
diff --git a/tools/lab/setup.py b/tools/lab/setup.py
index c2c77baaab..20e9daafb3 100755
--- a/tools/lab/setup.py
+++ b/tools/lab/setup.py
@@ -74,7 +74,7 @@ class LabHealthInstallDependencies(cmd.Command):
 def main():
     setuptools.setup(
         name='LabHealth',
-        version='0.1',
+        version='0.3',
         description='Android Test Lab Health',
         license='Apache2.0',
         cmdclass={
diff --git a/tools/lab/tests/fake.py b/tools/lab/tests/fake.py
index 07aa4e0dcd..be22e2ad5f 100644
--- a/tools/lab/tests/fake.py
+++ b/tools/lab/tests/fake.py
@@ -57,7 +57,7 @@ class MockShellCommand(object):
         else:
             return self._fake_result
 
-    def get_pids(self, identifier):
+    def get_command_pids(self, identifier):
         """Returns a generator of fake pids
 
         Args:
diff --git a/tools/lab/utils/shell.py b/tools/lab/utils/shell.py
index d1e8c484de..296fda1445 100644
--- a/tools/lab/utils/shell.py
+++ b/tools/lab/utils/shell.py
@@ -91,6 +91,38 @@ class ShellCommand(object):
         except job.Error:
             return False
 
+    def get_command_pids(self, identifier):
+        """Gets the pids of a program, based only upon the starting command
+
+        Searches for a program with a specific name and grabs the pids for all
+        programs that match only the command that started it. The arguments of
+        the command are ignored.
+
+        Args:
+          identifier: The search term that identifies the program.
+
+        Returns:
+          An array of ints of all pids that matched the identifier, or None if
+          no pids were found.
+
+        Raises:
+            StopIteration upon unsuccessful run command.
+        """
+        try:
+            result = self.run('ps -C %s --no-heading -o pid:1' % identifier)
+        except job.Error:
+            raise StopIteration
+
+        # Output looks like pids on separate lines:
+        # 1245
+        # 5001
+        pids = result.stdout.splitlines()
+
+        if pids:
+            return [int(pid) for pid in result.stdout.splitlines()]
+        else:
+            return None
+
     def get_pids(self, identifier):
         """Gets the pids of a program.