cros_interface: add a GetFile method

Adds a method to copy a file from the device to a particular local
filename and use it as part of GetFileContents, and add unit tests.

BUG=chromium:266400
TEST=CrOSInterface unit tests pass

Review URL: https://codereview.chromium.org/196833008

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@259068 0039d316-1c4b-4281-b951-d872f2087c98

2 files changed, 76 insertions, 20 deletions

diff --git a/tools/telemetry/telemetry/core/backends/chrome/cros_interface.py b/tools/telemetry/telemetry/core/backends/chrome/cros_interface.py
index 1838e7f..bb49ced 100644
--- a/tools/telemetry/telemetry/core/backends/chrome/cros_interface.py
+++ b/tools/telemetry/telemetry/core/backends/chrome/cros_interface.py
@@ -5,6 +5,7 @@
 import logging
 import os
 import re
+import shutil
 import subprocess
 import sys
 import tempfile
@@ -221,26 +222,52 @@ class CrOSInterface(object):
       f.flush()
       self.PushFile(f.name, remote_filename)
 
-  def GetFileContents(self, filename):
-    assert not self.local
-    with tempfile.NamedTemporaryFile() as f:
-      args = ['scp'] + self._ssh_args
-      if self._ssh_identity:
-        args.extend(['-i', self._ssh_identity])
+  def GetFile(self, filename, destfile=None):
+    """Copies a local file |filename| to |destfile| on the device.
 
-      args.extend(['root@%s:%s' % (self._hostname, filename),
-                   os.path.abspath(f.name)])
+    Args:
+      filename: The name of the local source file.
+      destfile: The name of the file to copy to, and if it is not specified
+        then it is the basename of the source file.
 
-      stdout, stderr = GetAllCmdOutput(args, quiet=True)
-      stderr = self._RemoveSSHWarnings(stderr)
+    """
+    logging.debug("GetFile(%s, %s)" % (filename, destfile))
+    if self.local:
+      if destfile is not None and destfile != filename:
+        shutil.copyfile(filename, destfile)
+      return
 
-      if stderr != '':
-        raise OSError('No such file or directory %s' % stderr)
+    if destfile is None:
+      destfile = os.path.basename(filename)
+    args = ['scp'] + self._ssh_args
+    if self._ssh_identity:
+      args.extend(['-i', self._ssh_identity])
 
-      with open(f.name, 'r') as f2:
-        res = f2.read()
-        logging.debug("GetFileContents(%s)->%s" % (filename, res))
-        return res
+    args.extend(['root@%s:%s' % (self._hostname, filename),
+                 os.path.abspath(destfile)])
+    stdout, stderr = GetAllCmdOutput(args, quiet=True)
+    stderr = self._RemoveSSHWarnings(stderr)
+    if stderr != '':
+      raise OSError('No such file or directory %s' % stderr)
+
+  def GetFileContents(self, filename):
+    """Get the contents of a file on the device.
+
+    Args:
+      filename: The name of the file on the device.
+
+    Returns:
+      A string containing the contents of the file.
+    """
+    # TODO: handle the self.local case
+    assert not self.local
+    t = tempfile.NamedTemporaryFile()
+    self.GetFile(filename, t.name)
+    with open(t.name, 'r') as f2:
+      res = f2.read()
+      logging.debug("GetFileContents(%s)->%s" % (filename, res))
+      f2.close()
+      return res
 
   def ListProcesses(self):
     """Returns (pid, cmd, ppid, state) of all processes on the device."""
diff --git a/tools/telemetry/telemetry/core/backends/chrome/cros_interface_unittest.py b/tools/telemetry/telemetry/core/backends/chrome/cros_interface_unittest.py
index 5b75a6a..896cb6b 100644
--- a/tools/telemetry/telemetry/core/backends/chrome/cros_interface_unittest.py
+++ b/tools/telemetry/telemetry/core/backends/chrome/cros_interface_unittest.py
@@ -7,6 +7,7 @@
 # a bit.
 
 import socket
+import tempfile
 import unittest
 
 from telemetry import test
@@ -52,18 +53,46 @@ class CrOSInterfaceTest(unittest.TestCase):
     cri = cros_interface.CrOSInterface(
       remote,
       options_for_unittests.GetCopy().cros_ssh_identity)
-    hosts = cri.GetFileContents('/etc/hosts')
-    assert hosts.startswith('# /etc/hosts')
+    hosts = cri.GetFileContents('/etc/lsb-release')
+    self.assertTrue('CHROMEOS' in hosts)
 
   @test.Enabled('cros-chrome')
-  def testGetFileContentsForSomethingThatDoesntExist(self):
+  def testGetFileContentsNonExistent(self):
     remote = options_for_unittests.GetCopy().cros_remote
     cri = cros_interface.CrOSInterface(
       remote,
       options_for_unittests.GetCopy().cros_ssh_identity)
+    f = tempfile.NamedTemporaryFile()
+    cri.PushContents('testGetFileNonExistent', f.name)
+    cri.RmRF(f.name)
     self.assertRaises(
       OSError,
-      lambda: cri.GetFileContents('/tmp/209fuslfskjf/dfsfsf'))
+      lambda: cri.GetFileContents(f.name))
+
+  @test.Enabled('cros-chrome')
+  def testGetFile(self): # pylint: disable=R0201
+    remote = options_for_unittests.GetCopy().cros_remote
+    cri = cros_interface.CrOSInterface(
+      remote,
+      options_for_unittests.GetCopy().cros_ssh_identity)
+    f = tempfile.NamedTemporaryFile()
+    cri.GetFile('/etc/lsb-release', f.name)
+    with open(f.name, 'r') as f2:
+      res = f2.read()
+      self.assertTrue('CHROMEOS' in res)
+
+  @test.Enabled('cros-chrome')
+  def testGetFileNonExistent(self):
+    remote = options_for_unittests.GetCopy().cros_remote
+    cri = cros_interface.CrOSInterface(
+      remote,
+      options_for_unittests.GetCopy().cros_ssh_identity)
+    f = tempfile.NamedTemporaryFile()
+    cri.PushContents('testGetFileNonExistent', f.name)
+    cri.RmRF(f.name)
+    self.assertRaises(
+      OSError,
+      lambda: cri.GetFile(f.name))
 
   @test.Enabled('cros-chrome')
   def testIsServiceRunning(self):