Don't patch non-stub drives

Signed-off-by: Kazuyoshi Kato <[email protected]>

Author: kzys

runtime/drive_handler.go

@@ -125,10 +125,20 @@ func (h *stubDriveHandler) createStubDrive(driveID, path string) error {
 	return nil
 }
 
+func filterOnlyStubDrives(drives []models.Drive) []models.Drive {
+	stubDrives := make([]models.Drive, 0)
+	for _, d := range drives {
+		if !*d.IsRootDevice {
+			stubDrives = append(stubDrives, d)
+		}
+	}
+	return stubDrives
+}
+
 // SetDrives will set the given drives and the offset to which the stub drives
 // start.
 func (h *stubDriveHandler) SetDrives(offset int64, d []models.Drive) {
-	h.drives = d
+	h.drives = filterOnlyStubDrives(d)
 	h.stubDriveIndex = offset
 }
 

runtime/drive_handler_test.go

@@ -15,6 +15,7 @@ package main
 
 import (
 	"context"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -102,6 +103,61 @@ func TestPatchStubDrive(t *testing.T) {
 	}
 }
 
+func TestPatchStubDrive_single(t *testing.T) {
+	ctx := context.Background()
+	index := 0
+	expectedReplacements := []string{
+		"/correct/path0",
+		"/correct/path1",
+	}
+
+	mockClient := &fctesting.MockClient{
+		PatchGuestDriveByIDFn: func(params *ops.PatchGuestDriveByIDParams) (*ops.PatchGuestDriveByIDNoContent, error) {
+			assert.Equal(t, expectedReplacements[index], firecracker.StringValue(params.Body.PathOnHost))
+			index++
+			return nil, nil
+		},
+	}
+
+	logger := log.G(ctx)
+
+	fcClient := firecracker.NewClient("/path/to/socket", nil, false, firecracker.WithOpsClient(mockClient))
+	client, err := firecracker.NewMachine(ctx, firecracker.Config{}, firecracker.WithClient(fcClient))
+	assert.NoError(t, err, "failed to create new machine")
+
+	tempDir := os.TempDir()
+	path, err := ioutil.TempDir(tempDir, t.Name())
+	assert.NoError(t, err, "failed to create test directory")
+	defer os.RemoveAll(path)
+
+	handler := newStubDriveHandler(path, logger)
+
+	driveBuilder := firecracker.NewDrivesBuilder("fc/root")
+
+	stubDriveIndex := int64(len(driveBuilder.Build()) - 1)
+	containerCount := 1
+
+	paths, err := handler.StubDrivePaths(containerCount)
+	assert.NoError(t, err)
+
+	for i, path := range paths {
+		driveID := fmt.Sprintf("stub%d", i)
+		driveBuilder = driveBuilder.AddDrive(path, false, func(drive *models.Drive) {
+			drive.DriveID = firecracker.String(driveID)
+		})
+	}
+
+	handler.SetDrives(stubDriveIndex, driveBuilder.Build())
+
+	drive1, err := handler.PatchStubDrive(ctx, client, "/correct/path0")
+	assert.NoError(t, err)
+	assert.Equal(t, "stub0", firecracker.StringValue(drive1))
+
+	drive2, err := handler.PatchStubDrive(ctx, client, "/correct/path1")
+	assert.NotEqual(t, "root_drive", firecracker.StringValue(drive2))
+	assert.Error(t, err)
+}
+
 func TestPatchStubDrive_concurrency(t *testing.T) {
 	ctx := context.Background()
 	mockClient := &fctesting.MockClient{