Merge pull request #9083 from thaJeztah/lease_withlabel

leases: add WithLabel

cmd/ctr/commands/images/mount.go

@@ -75,9 +75,7 @@ When you are done, use the unmount command.
 		ctx, done, err := client.WithLease(ctx,
 			leases.WithID(target),
 			leases.WithExpiration(24*time.Hour),
-			leases.WithLabels(map[string]string{
+			leases.WithLabel("containerd.io/gc.ref.snapshot."+snapshotter, target),
-				"containerd.io/gc.ref.snapshot." + snapshotter: target,
-			}),
 		)
 		if err != nil && !errdefs.IsAlreadyExists(err) {
 			return err

leases/lease.go

@@ -65,6 +65,19 @@ func SynchronousDelete(ctx context.Context, o *DeleteOptions) error {
 	return nil
 }
 
+// WithLabel sets a label on a lease, and merges it with existing labels.
+// It overwrites the existing value of the given label (if present).
+func WithLabel(label, value string) Opt {
+	return func(l *Lease) error {
+		if l.Labels == nil {
+			l.Labels = map[string]string{label: value}
+			return nil
+		}
+		l.Labels[label] = value
+		return nil
+	}
+}
+
 // WithLabels merges labels on a lease
 func WithLabels(labels map[string]string) Opt {
 	return func(l *Lease) error {

leases/lease_test.go

@@ -24,31 +24,26 @@ import (
 )
 
 func TestWithLabels(t *testing.T) {
-	type unitTest struct {
+	testcases := []struct {
 		name          string
-		uut      *Lease
+		initialLabels map[string]string
 		labels        map[string]string
 		expected      map[string]string
-	}
+	}{
-
+		{
-	addLabelsToEmptyMap := &unitTest{
 			name: "AddLabelsToEmptyMap",
-		uut:  &Lease{},
 			labels: map[string]string{
 				"containerd.io/gc.root": "2015-12-04T00:00:00Z",
 			},
 			expected: map[string]string{
 				"containerd.io/gc.root": "2015-12-04T00:00:00Z",
 			},
-	}
+		},
-
+		{
-	addLabelsToNonEmptyMap := &unitTest{
 			name: "AddLabelsToNonEmptyMap",
-		uut: &Lease{
+			initialLabels: map[string]string{
-			Labels: map[string]string{
 				"containerd.io/gc.expire": "2015-12-05T00:00:00Z",
 			},
-		},
 			labels: map[string]string{
 				"containerd.io/gc.root":                   "2015-12-04T00:00:00Z",
 				"containerd.io/gc.ref.snapshot.overlayfs": "sha256:87806a591ce894ff5c699c28fe02093d6cdadd6b1ad86819acea05ccb212ff3d",
@@ -58,26 +53,38 @@ func TestWithLabels(t *testing.T) {
 				"containerd.io/gc.ref.snapshot.overlayfs": "sha256:87806a591ce894ff5c699c28fe02093d6cdadd6b1ad86819acea05ccb212ff3d",
 				"containerd.io/gc.expire":                 "2015-12-05T00:00:00Z",
 			},
+		},
 	}
 
-	testcases := []*unitTest{
+	for _, tc := range testcases {
-		addLabelsToEmptyMap,
+		tc := tc
-		addLabelsToNonEmptyMap,
+		t.Run(tc.name, func(t *testing.T) {
-	}
+			lease := newLease(tc.initialLabels)
-
+			err := WithLabels(tc.labels)(lease)
-	for _, testcase := range testcases {
-		testcase := testcase
-
-		t.Run(testcase.name, func(t *testing.T) {
-			f := WithLabels(testcase.labels)
-
-			err := f(testcase.uut)
 			require.NoError(t, err)
-
+			assert.Equal(t, lease.Labels, tc.expected)
-			for k, v := range testcase.expected {
+		})
-				assert.Contains(t, testcase.uut.Labels, k)
-				assert.Equal(t, v, testcase.uut.Labels[k])
 	}
+	for _, tc := range testcases {
+		tc := tc
+		t.Run(tc.name+"-WithLabel", func(t *testing.T) {
+			lease := newLease(tc.initialLabels)
+			for k, v := range tc.labels {
+				err := WithLabel(k, v)(lease)
+				require.NoError(t, err)
+			}
+			assert.Equal(t, lease.Labels, tc.expected)
 		})
 	}
 }
+
+func newLease(labels map[string]string) *Lease {
+	lease := &Lease{}
+	if labels != nil {
+		lease.Labels = map[string]string{}
+		for k, v := range labels {
+			lease.Labels[k] = v
+		}
+	}
+	return lease
+}