Finding main device in test

internal/storage/storage.go

@@ -83,7 +83,9 @@ func GetDevices() []Device {
 
 	mainDiskID := "sda"
 
-	if diagnostics.Version == "prod" || diagnostics.Version == "dev" {
+	if diagnostics.Version == "dev" {
+		mainDiskID = "sda"
+	} else if diagnostics.Version == "prod" {
 		mainDiskID = "mmcblk0"
 	} else if diagnostics.Version == "cloud" {
 		mainDiskID = "vda"

internal/storage/storage_test.go

@@ -14,17 +14,25 @@ func TestGetDevices(t *testing.T) {
 		t.Fail()
 	}
 
+	foundMainDevice := false
 	foundDevice := false
 
 	t.Log("Looking for a mmcblk0 or sda device")
 	for _, device := range result {
+
+		if device.MainDevice {
+			t.Log("Found target main device", device.ID)
+			foundMainDevice = true
+		}
+
 		if device.ID == "mmcblk0" || device.ID == "sda" {
 			t.Log("Found target device", device.ID)
 			foundDevice = true
 		}
+
 	}
 
-	if !foundDevice {
+	if !foundDevice || !foundMainDevice {
 		t.Log("Expected to find device mmcblk0 but did not. Devices:", result)
 		t.Fail()
 	}