Fix TestCreateVM_Isolated stability by creating fresh ctx/fcClient per subtest

Signed-off-by: Arjun Raja Yogidas <arjunry@amazon.com>

Author: coderbirju

runtime/service_integ_test.go

@@ -112,7 +112,9 @@ func iperf3Image(ctx context.Context, client *containerd.Client, snapshotterName
 func TestShimExitsUponContainerDelete_Isolated(t *testing.T) {
 	integtest.Prepare(t)
 
-	ctx := namespaces.WithNamespace(context.Background(), defaultNamespace)
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
 
 	client, err := containerd.New(integtest.ContainerdSockPath)
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
@@ -732,7 +734,9 @@ func TestLongUnixSocketPath_Isolated(t *testing.T) {
 	namespace := strings.Repeat("n", 20)
 	vmID := strings.Repeat("v", 64)
 
-	ctx := namespaces.WithNamespace(context.Background(), namespace)
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, namespace)
 
 	fcClient, err := integtest.NewFCControlClient(integtest.ContainerdSockPath)
 	require.NoError(t, err, "failed to create fccontrol client")
@@ -827,7 +831,9 @@ func TestStubBlockDevices_Isolated(t *testing.T) {
 
 	const vmID = 0
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	client, err := containerd.New(integtest.ContainerdSockPath, containerd.WithDefaultRuntime(firecrackerRuntime))
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
@@ -963,7 +969,9 @@ func startAndWaitTask(ctx context.Context, t testing.TB, c containerd.Container)
 }
 
 func testCreateContainerWithSameName(t *testing.T, vmID string) {
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	// Explicitly specify Container Count = 2 to workaround #230
 	if len(vmID) != 0 {
@@ -1051,7 +1059,9 @@ func TestCreateContainerWithSameName_Isolated(t *testing.T) {
 func TestStubDriveReserveAndReleaseByContainers_Isolated(t *testing.T) {
 	integtest.Prepare(t)
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	client, err := containerd.New(integtest.ContainerdSockPath, containerd.WithDefaultRuntime(firecrackerRuntime))
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
@@ -1439,7 +1449,9 @@ func TestUpdateVMMetadata_Isolated(t *testing.T) {
 
 func TestMemoryBalloon_Isolated(t *testing.T) {
 	integtest.Prepare(t)
-	ctx := namespaces.WithNamespace(context.Background(), defaultNamespace)
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, defaultNamespace)
 
 	numberOfVms := defaultNumberOfVms
 	if str := os.Getenv(numberOfVmsEnvName); str != "" {
@@ -1680,7 +1692,9 @@ func TestStopVM_Isolated(t *testing.T) {
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
 	defer client.Close()
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	image, err := alpineImage(ctx, client, defaultSnapshotterName)
 	require.NoError(t, err, "failed to get alpine image")
@@ -1905,7 +1919,9 @@ func TestStopVMFailFast_Isolated(t *testing.T) {
 
 	name := testNameToVMID(t.Name())
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	image, err := alpineImage(ctx, client, defaultSnapshotterName)
 	require.NoError(t, err, "failed to get alpine image")
@@ -1957,7 +1973,9 @@ func TestExec_Isolated(t *testing.T) {
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
 	defer client.Close()
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	image, err := alpineImage(ctx, client, defaultSnapshotterName)
 	require.NoError(t, err, "failed to get alpine image")
@@ -2073,7 +2091,9 @@ func TestEvents_Isolated(t *testing.T) {
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
 	defer client.Close()
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	// If we don't have enough events within 30 seconds, the context will be cancelled and the loop below will be interrupted
 	subscribeCtx, subscribeCancel := context.WithTimeout(ctx, 30*time.Second)
@@ -2159,7 +2179,9 @@ func TestOOM_Isolated(t *testing.T) {
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
 	defer client.Close()
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	// If we don't have enough events within 30 seconds, the context will be cancelled and the loop below will be interrupted
 	subscribeCtx, subscribeCancel := context.WithTimeout(ctx, 30*time.Second)
@@ -2250,17 +2272,12 @@ func TestCreateVM_Isolated(t *testing.T) {
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
 	defer client.Close()
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
-
-	fcClient, err := integtest.NewFCControlClient(integtest.ContainerdSockPath)
-	require.NoError(t, err, "failed to create ttrpc client")
-
 	kernelArgs := integtest.DefaultRuntimeConfig.KernelArgs
 
 	type subtest struct {
 		name                    string
 		request                 proto.CreateVMRequest
-		validate                func(*testing.T, error)
+		validate                func(*testing.T, context.Context, error)
 		validateUsesFindProcess bool
 		stopVM                  bool
 	}
@@ -2269,7 +2286,7 @@ func TestCreateVM_Isolated(t *testing.T) {
 		{
 			name:    "Happy Case",
 			request: proto.CreateVMRequest{},
-			validate: func(t *testing.T, err error) {
+			validate: func(t *testing.T, ctx context.Context, err error) {
 				require.NoError(t, err)
 			},
 			stopVM: true,
@@ -2283,7 +2300,7 @@ func TestCreateVM_Isolated(t *testing.T) {
 					HostPath: "/var/lib/firecracker-containerd/runtime/rootfs-debug.img",
 				},
 			},
-			validate: func(t *testing.T, err error) {
+			validate: func(t *testing.T, ctx context.Context, err error) {
 				require.NotNil(t, err, "expected an error but did not receive any")
 				time.Sleep(5 * time.Second)
 				firecrackerProcesses, err := findProcess(ctx, findFirecracker)
@@ -2300,11 +2317,12 @@ func TestCreateVM_Isolated(t *testing.T) {
 				RootDrive: &proto.FirecrackerRootDrive{
 					HostPath: "/var/lib/firecracker-containerd/runtime/rootfs-debug.img",
 				},
+				TimeoutSeconds: 5,
 			},
-			validate: func(t *testing.T, err error) {
+			validate: func(t *testing.T, ctx context.Context, err error) {
 				require.Error(t, err)
 				assert.Equal(t, codes.DeadlineExceeded, status.Code(err))
-				assert.Contains(t, err.Error(), "didn't start within 20s")
+				assert.Contains(t, err.Error(), "didn't start within 5s")
 			},
 			stopVM: true,
 		},
@@ -2317,7 +2335,7 @@ func TestCreateVM_Isolated(t *testing.T) {
 				},
 				TimeoutSeconds: 60,
 			},
-			validate: func(t *testing.T, err error) {
+			validate: func(t *testing.T, ctx context.Context, err error) {
 				require.NoError(t, err)
 			},
 			stopVM: true,
@@ -2331,6 +2349,15 @@ func TestCreateVM_Isolated(t *testing.T) {
 		if !s.validateUsesFindProcess {
 			t.Parallel()
 		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+		defer cancel()
+		ctx = namespaces.WithNamespace(ctx, "default")
+
+		fcClient, err := integtest.NewFCControlClient(integtest.ContainerdSockPath)
+		require.NoError(t, err, "failed to create ttrpc client")
+		defer fcClient.Close()
+
 		vmID := testNameToVMID(t.Name())
 
 		tempDir := t.TempDir()
@@ -2349,7 +2376,7 @@ func TestCreateVM_Isolated(t *testing.T) {
 		requireNonEmptyFifo(t, metricsFile)
 
 		// Some test cases are expected to have an error, some are not.
-		s.validate(t, createVMErr)
+		s.validate(t, ctx, createVMErr)
 
 		if createVMErr == nil && s.stopVM {
 			// Ensure the response fields are populated correctly
@@ -2360,6 +2387,8 @@ func TestCreateVM_Isolated(t *testing.T) {
 				TimeoutSeconds: 30,
 			})
 			require.Equal(t, status.Code(err), codes.OK)
+
+			time.Sleep(500 * time.Millisecond)
 		}
 	}
 
@@ -2525,7 +2554,9 @@ func TestAttach_Isolated(t *testing.T) {
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
 	defer client.Close()
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	image, err := alpineImage(ctx, client, defaultSnapshotterName)
 	require.NoError(t, err, "failed to get alpine image")
@@ -2627,7 +2658,9 @@ func TestBrokenPipe_Isolated(t *testing.T) {
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", integtest.ContainerdSockPath)
 	defer client.Close()
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	ctx = namespaces.WithNamespace(ctx, "default")
 
 	image, err := alpineImage(ctx, client, defaultSnapshotterName)
 	require.NoError(t, err, "failed to get alpine image")
@@ -2641,6 +2674,7 @@ func TestBrokenPipe_Isolated(t *testing.T) {
 		containerd.WithNewSpec(oci.WithProcessArgs("/usr/bin/yes")),
 	)
 	require.NoError(t, err)
+	defer c.Delete(ctx, containerd.WithSnapshotCleanup)
 
 	var stdout1 errorBuffer
 	var stderr1 errorBuffer