diff --git a/requirements.txt b/requirements.txt index b0528704f..e43dbace9 100644 --- a/requirements.txt +++ b/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 diff --git a/teeth_agent/hardware.py b/teeth_agent/hardware.py index e4dd68bab..095065c71 100644 --- a/teeth_agent/hardware.py +++ b/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): diff --git a/teeth_agent/tests/hardware.py b/teeth_agent/tests/hardware.py index ca2d89bfd..bc0902fb7 100644 --- a/teeth_agent/tests/hardware.py +++ b/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')