utils: Add partition concatenation routine.

Follow kernel convention and add 'p' if the block device ends with a digit.
Fixes platform_PartitionCheck test.
Follow same method for faft test.

BUG=chromium:776029
TEST=unit test, test latform_PartitionCheck on NVMe.

Change-Id: I61f2f194e149634e132a4ad548335ea7c10d8366
Signed-off-by: Gwendal Grignou <gwendal@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/741483
Reviewed-by: Sridhar Sonti <sontis@google.com>
Reviewed-by: Wai-Hong Tam <waihong@google.com>
diff --git a/client/bin/unit_test.py b/client/bin/unit_test.py
index 3231daf..55a8c9a 100644
--- a/client/bin/unit_test.py
+++ b/client/bin/unit_test.py
@@ -19,29 +19,29 @@
       os.chdir(self.srcdir)
       utils.make('clean')
       utils.make('all')
-      
+
       self.job.setup_dep(['gtest'])
 
     def run_once(self):
         dep ='gtest'
         dep_dir = os.path.join(self.autodir, 'deps', dep)
         self.job.install_pkg(dep, 'dep', dep_dir)
-      
+
         # Run the unit test, gather the results and place the gcda files for
         # code coverage in the results directory.
-        
+
         os.chdir(self.srcdir)
-        result = utils.run('LD_LIBRARY_PATH=' + dep_dir + 
+        result = utils.run('LD_LIBRARY_PATH=' + dep_dir +
                            ' GCOV_PREFIX=' + self.resultsdir +
                            ' GCOV_PREFIX_STRIP=9999 ./unit_test > ' +
                            self.resultsdir + '/unit_test_result.txt')
         logging.debug(result.stderr)
         logging.info('result: ' + self.resultsdir + '/unit_test_result.txt')
-        
+
     def cleanup(self):
         # This is a hack - we should only need to copy back the .gcda file but
-        # we don't know how to access the source on the server. So copy 
+        # we don't know how to access the source on the server. So copy
         # everything back.
-        
+
         os.chdir(self.srcdir)
         utils.run('cp * ' + self.resultsdir)
diff --git a/client/bin/utils.py b/client/bin/utils.py
index 79b4d15..a60b658 100644
--- a/client/bin/utils.py
+++ b/client/bin/utils.py
@@ -910,6 +910,20 @@
     # For harddisk rtt > 0
     return rtt and int(rtt) > 0
 
+def concat_partition(disk_name, partition_number):
+    """
+    Return the name of a partition:
+    sda, 3 --> sda3
+    mmcblk0, 3 --> mmcblk0p3
+
+    @param disk_name: diskname string
+    @param partition_number: integer
+    """
+    if disk_name.endswith(tuple(str(i) for i in range(0, 10))):
+        sep = 'p'
+    else:
+        sep = ''
+    return disk_name + sep + str(partition_number)
 
 def verify_hdparm_feature(disk_name, feature):
     """
diff --git a/client/bin/utils_unittest.py b/client/bin/utils_unittest.py
new file mode 100755
index 0000000..d89e178
--- /dev/null
+++ b/client/bin/utils_unittest.py
@@ -0,0 +1,19 @@
+#!/usr/bin/python
+
+__author__ = "gwendal@google.com (Gwendal Grignou)"
+
+import unittest
+
+from autotest_lib.client.bin import utils
+
+class TestUtils(unittest.TestCase):
+    """Test utils functions."""
+
+
+    def test_concat_partition(self):
+        self.assertEquals("nvme0n1p3", utils.concat_partition("nvme0n1", 3))
+        self.assertEquals("mmcblk1p3", utils.concat_partition("mmcblk1", 3))
+        self.assertEquals("sda3", utils.concat_partition("sda", 3))
+
+
+
diff --git a/client/cros/faft/utils/os_interface.py b/client/cros/faft/utils/os_interface.py
index f2008cf..40ea188 100644
--- a/client/cros/faft/utils/os_interface.py
+++ b/client/cros/faft/utils/os_interface.py
@@ -289,9 +289,7 @@
 
     def join_part(self, dev, part):
         """Return a concatenated string of device and partition number"""
-        if 'mmcblk' in dev:
-            return dev + 'p' + part
-        if 'nvme' in dev:
+        if dev.endswith(tuple(str(i) for i in range(0, 10))):
             return dev + 'p' + part
         else:
             return dev + part
diff --git a/client/site_tests/platform_PartitionCheck/platform_PartitionCheck.py b/client/site_tests/platform_PartitionCheck/platform_PartitionCheck.py
index 1a5cded..ed1722e 100644
--- a/client/site_tests/platform_PartitionCheck/platform_PartitionCheck.py
+++ b/client/site_tests/platform_PartitionCheck/platform_PartitionCheck.py
@@ -53,21 +53,15 @@
     def run_once(self):
         errors = []
         device = os.path.basename(utils.get_fixed_dst_drive())
-        mmcpath = os.path.join('/sys', 'block', device)
-
-        if os.path.exists(mmcpath) and device.startswith('mmc'):
-            partitions = [device + 'p3', device + 'p5']
-        else:
-            partitions = [device + '3', device + '5']
-
+        partitions = [utils.concat_partition(device, i) for i in (3, 5)]
         block_size = self.get_block_size(device)
 
         for p in partitions:
             pblocks = self.get_partition_size(device, p)
             psize = pblocks * block_size
             if psize != ROOTFS_SIZE_2G and psize != ROOTFS_SIZE_4G:
-                errmsg = ('%s is %d bytes, expected %d' %
-                          (p, psize, ROOTFS_SIZE))
+                errmsg = ('%s is %d bytes, expected %d or %d' %
+                          (p, psize, ROOTFS_SIZE_2G, ROOTFS_SIZE_4G))
                 logging.warning(errmsg)
                 errors.append(errmsg)