From f84b9527b9000dec13172e66f72f6d1934f17d93 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <flamingcow@gmail.com>
Date: Thu, 31 Mar 2016 17:01:50 -0700
Subject: [PATCH] Handle nvme drives

---
 server/image.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/server/image.py b/server/image.py
index 2f2632c..afe67cb 100755
--- a/server/image.py
+++ b/server/image.py
@@ -6,7 +6,6 @@ import shutil
 import subprocess
 import sys
 import tempfile
-import time
 
 
 parser = argparse.ArgumentParser(description='iconograph build_image')
@@ -55,8 +54,18 @@ class Imager(object):
     subprocess.check_call(args, **kwargs)
 
   def _PartDev(self, part_num):
-    # TODO: better way to figure out what the partition device is
-    return '%s%d' % (self._device, part_num)
+    args = {
+      'device': self._device,
+      'part': part_num,
+    }
+    options = [
+      '%(device)s%(part)d' % args,
+      '%(device)sp%(part)d' % args,
+    ]
+    while True:
+      for option in options:
+        if os.path.exists(option):
+          return option
 
   def _PartitionAndMkFS(self):
     self._Exec(
@@ -72,7 +81,6 @@ class Imager(object):
         '--align', 'optimal',
         self._device,
         'mkpart', 'primary', 'ext4', '0%', boot_stop)
-    time.sleep(1) # yuck
     self._Exec(
         'mkfs.ext4',
         '-L', 'BOOT',
@@ -86,7 +94,6 @@ class Imager(object):
           '--align', 'optimal',
           self._device,
           'mkpart', 'primary', 'ext4', boot_stop, '100%')
-      time.sleep(1) # yuck
       self._Exec(
           'mkfs.ext4',
           '-L', 'PERSISTENT',