Merge pull request #38 from rackerlabs/detect_install_disk

Detect OS Installation Disk

requirements.txt

@@ -1,5 +1,6 @@
 Werkzeug==0.9.4
 requests==2.0.0
 cherrypy==3.2.4
+plumbum==1.4.0
 stevedore==0.13
 -e git+https://github.com/racker/teeth-rest.git@e876c0fddd5ce2f5223ab16936f711b0d57e19c4#egg=teeth_rest

teeth_agent/hardware.py

@@ -16,6 +16,7 @@ limitations under the License.
 
 import abc
 
+import plumbum
 import stevedore
 import structlog
 
@@ -35,6 +36,13 @@ class HardwareSupport(object):
     SERVICE_PROVIDER = 3
 
 
+class BlockDevice(object):
+    def __init__(self, name, size, start_sector):
+        self.name = name
+        self.size = size
+        self.start_sector = start_sector
+
+
 class HardwareManager(object):
     @abc.abstractmethod
     def evaluate_hardware_support(cls):
@@ -56,8 +64,35 @@ class GenericHardwareManager(HardwareManager):
     def get_primary_mac_address(self):
         return open('/sys/class/net/eth0/address', 'r').read().strip('\n')
 
+    def _cmd(self, command_name):
+        """Mocking plumbum is frustratingly difficult. Instead, mock this."""
+        return plumbum.local[command_name]
+
+    def _list_block_devices(self):
+        report = self._cmd('blockdev')('--report').strip()
+        lines = report.split('\n')
+        lines = [line.split() for line in lines]
+        startsec_idx = lines[0].index('StartSec')
+        device_idx = lines[0].index('Device')
+        size_idx = lines[0].index('Size')
+        return [BlockDevice(line[device_idx],
+                            int(line[size_idx]),
+                            int(line[startsec_idx]))
+                for line
+                in lines[1:]]
+
     def get_os_install_device(self):
-        return '/dev/sda'
+        # Assume anything with a start sector other than 0, is a partition
+        block_devices = [device for device in self._list_block_devices()
+                         if device.start_sector == 0]
+
+        # Find the first device larger than 4GB, assume it is the OS disk
+        # TODO(russellhaering): This isn't a valid assumption in all cases,
+        #                       is there a more reasonable default behavior?
+        block_devices.sort(key=lambda device: device.size)
+        for device in block_devices:
+            if device.size >= (4 * pow(1024, 3)):
+                return device.name
 
 
 def _compare_extensions(ext1, ext2):

teeth_agent/tests/hardware.py

@@ -36,4 +36,16 @@ class TestGenericHardwareManager(unittest.TestCase):
         f.read.assert_called_once_with()
 
     def test_get_os_install_device(self):
-        self.assertEqual(self.hardware.get_os_install_device(), '/dev/sda')
+        self.hardware._cmd = mock.Mock()
+        self.hardware._cmd.return_value = blockdev = mock.Mock()
+        blockdev.return_value = (
+            'RO    RA   SSZ   BSZ   StartSec            Size   Device\n'
+            'rw   256   512  4096          0    249578283616   /dev/sda\n'
+            'rw   256   512  4096       2048      8587837440   /dev/sda1\n'
+            'rw   256   512  4096  124967424        15728640   /dev/sda2\n'
+            'rw   256   512  4096          0     31016853504   /dev/sdb\n'
+            'rw   256   512  4096          0    249578283616   /dev/sdc\n')
+
+        self.assertEqual(self.hardware.get_os_install_device(), '/dev/sdb')
+        self.hardware._cmd.assert_called_once_with('blockdev')
+        blockdev.assert_called_once_with('--report')