diff --git a/tests/assets/test_01/cloud-config.yml b/tests/assets/test_01/cloud-config.yml
index 1b5de81d..d3fa591b 100644
--- a/tests/assets/test_01/cloud-config.yml
+++ b/tests/assets/test_01/cloud-config.yml
@@ -3,8 +3,6 @@ rancher:
   environment:
     ETCD_DISCOVERY: https://discovery.etcd.io/c2c219023108cda9529364d6d983fe13
     FLANNEL_NETWORK: 10.244.0.0/16
-  services_include:
-    kernel-headers: true
   network:
     interfaces:
       eth*:
diff --git a/tests/assets/test_22/cloud-config.yml b/tests/assets/test_22/cloud-config.yml
new file mode 100644
index 00000000..abbd95f8
--- /dev/null
+++ b/tests/assets/test_22/cloud-config.yml
@@ -0,0 +1,6 @@
+#cloud-config
+rancher:
+  services_include:
+    kernel-headers: true
+ssh_authorized_keys:
+  - ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC85w9stZyiLQp/DkVO6fqwiShYcj1ClKdtCqgHtf+PLpJkFReSFu8y21y+ev09gsSMRRrjF7yt0pUHV6zncQhVeqsZtgc5WbELY2DOYUGmRn/CCvPbXovoBrQjSorqlBmpuPwsStYLr92Xn+VVsMNSUIegHY22DphGbDKG85vrKB8HxUxGIDxFBds/uE8FhSy+xsoyT/jUZDK6pgq2HnGl6D81ViIlKecpOpWlW3B+fea99ADNyZNVvDzbHE5pcI3VRw8u59WmpWOUgT6qacNVACl8GqpBvQk8sw7O/X9DSZHCKafeD9G5k+GYbAUz92fKWrx/lOXfUXPS3+c8dRIF
diff --git a/tests/common_test.go b/tests/common_test.go
index 03549e93..8abfa5e7 100644
--- a/tests/common_test.go
+++ b/tests/common_test.go
@@ -66,27 +66,6 @@ func (s *QemuSuite) RunQemu(additionalArgs ...string) error {
 	return s.WaitForSSH()
 }
 
-func (s *QemuSuite) RestartQemu(additionalArgs ...string) error {
-	s.qemuCmd.Process.Kill()
-	time.Sleep(time.Millisecond * 1000)
-
-	runArgs := []string{
-		"--qemu",
-		"--no-rebuild",
-		"--no-rm-usr",
-	}
-	runArgs = append(runArgs, additionalArgs...)
-
-	s.qemuCmd = exec.Command(s.runCommand, runArgs...)
-	s.qemuCmd.Stdout = os.Stdout
-	s.qemuCmd.Stderr = os.Stderr
-	if err := s.qemuCmd.Start(); err != nil {
-		return err
-	}
-
-	return s.WaitForSSH()
-}
-
 func (s *QemuSuite) WaitForSSH() error {
 	sshArgs := []string{
 		"--qemu",
@@ -124,9 +103,10 @@ func (s *QemuSuite) CheckCall(c *C, additionalArgs ...string) {
 	c.Assert(s.MakeCall(additionalArgs...), IsNil)
 }
 
-func (s *QemuSuite) Reboot() {
+func (s *QemuSuite) Reboot(c *C) {
 	s.MakeCall("sudo reboot")
 	time.Sleep(3000 * time.Millisecond)
+	c.Assert(s.WaitForSSH(), IsNil)
 }
 
 func (s *QemuSuite) LoadInstallerImage(c *C) {
diff --git a/tests/docker_in_persistent_console_test.go b/tests/docker_in_persistent_console_test.go
index 15920043..db4e61c0 100644
--- a/tests/docker_in_persistent_console_test.go
+++ b/tests/docker_in_persistent_console_test.go
@@ -14,7 +14,7 @@ func (s *QemuSuite) TestRebootWithContainerRunning(c *C) {
 set -e -x
 docker run -d --restart=always %s`, NginxImage))
 
-	s.Reboot()
+	s.Reboot(c)
 
 	s.CheckCall(c, "docker ps -f status=running | grep nginx")
 }
diff --git a/tests/http_proxy_test.go b/tests/http_proxy_test.go
index 761b46fd..af55747e 100644
--- a/tests/http_proxy_test.go
+++ b/tests/http_proxy_test.go
@@ -19,7 +19,7 @@ else
     exit 0
 fi`)
 
-	s.RestartQemu("--cloud-config", "./tests/assets/test_17/cloud-config.yml")
+	s.Reboot(c)
 
 	s.CheckCall(c, `
 set -x -e
diff --git a/tests/kernel_headers_test.go b/tests/kernel_headers_test.go
new file mode 100644
index 00000000..3297b85d
--- /dev/null
+++ b/tests/kernel_headers_test.go
@@ -0,0 +1,12 @@
+package integration
+
+import . "gopkg.in/check.v1"
+
+func (s *QemuSuite) TestKernelHeaders(c *C) {
+	err := s.RunQemu("--cloud-config", "./tests/assets/test_22/cloud-config.yml")
+	c.Assert(err, IsNil)
+
+	s.CheckCall(c, `
+sleep 10
+docker inspect kernel-headers`)
+}
diff --git a/tests/misc_test.go b/tests/misc_test.go
index e3f0f20a..49e58540 100644
--- a/tests/misc_test.go
+++ b/tests/misc_test.go
@@ -28,8 +28,4 @@ for i in $(pidof docker); do
     fi
 done
 [ "$found" = "true" ]`)
-
-	s.CheckCall(c, `
-sleep 5
-docker inspect kernel-headers`)
 }
diff --git a/tests/oem_test.go b/tests/oem_test.go
index 1ce23799..1095f390 100644
--- a/tests/oem_test.go
+++ b/tests/oem_test.go
@@ -3,14 +3,14 @@ package integration
 import . "gopkg.in/check.v1"
 
 func (s *QemuSuite) TestOem(c *C) {
-	err := s.RunQemu("--append", "rancher.state.dev=x")
+	err := s.RunQemu("--second-drive")
 	c.Assert(err, IsNil)
 
 	s.CheckCall(c, `
 set -x
 set -e
-sudo mkfs.ext4 -L RANCHER_OEM /dev/vda
-sudo mount /dev/vda /mnt
+sudo mkfs.ext4 -L RANCHER_OEM /dev/vdb
+sudo mount /dev/vdb /mnt
 cat > /tmp/oem-config.yml << EOF
 #cloud-config
 rancher:
@@ -20,10 +20,11 @@ EOF
 sudo cp /tmp/oem-config.yml /mnt
 sudo umount /mnt`)
 
-	s.Reboot()
+	s.Reboot(c)
 
 	s.CheckCall(c, `
 set -x
+set -e
 if [ ! -e /usr/share/ros/oem/oem-config.yml ]; then
     echo Failed to find /usr/share/ros/oem/oem-config.yml
     exit 1