From f7089ba225f83955222503034490d73231e62a3e Mon Sep 17 00:00:00 2001
From: Sebastiaan van Stijn <github@gone.nl>
Date: Mon, 11 Sep 2023 18:39:57 +0200
Subject: [PATCH] leases: add WithLabel

This adds a new WithLabel function, which allows to set a single label on
a lease, without having to first construct an intermediate map[string]string.

Signed-off-by: Sebastiaan van Stijn <github@gone.nl>
---
 cmd/ctr/commands/images/mount.go |  4 +--
 leases/lease.go                  | 13 ++++++++++
 leases/lease_test.go             | 42 +++++++++++++++++++++++---------
 3 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/cmd/ctr/commands/images/mount.go b/cmd/ctr/commands/images/mount.go
index 41b9e17bc..37a342da0 100644
--- a/cmd/ctr/commands/images/mount.go
+++ b/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{
-				"containerd.io/gc.ref.snapshot." + snapshotter: target,
-			}),
+			leases.WithLabel("containerd.io/gc.ref.snapshot."+snapshotter, target),
 		)
 		if err != nil && !errdefs.IsAlreadyExists(err) {
 			return err
diff --git a/leases/lease.go b/leases/lease.go
index fc0ca3491..d842c4081 100644
--- a/leases/lease.go
+++ b/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 {
diff --git a/leases/lease_test.go b/leases/lease_test.go
index c436f002b..786bcbdc0 100644
--- a/leases/lease_test.go
+++ b/leases/lease_test.go
@@ -25,14 +25,13 @@ import (
 
 func TestWithLabels(t *testing.T) {
 	testcases := []struct {
-		name     string
-		uut      *Lease
-		labels   map[string]string
-		expected map[string]string
+		name          string
+		initialLabels map[string]string
+		labels        map[string]string
+		expected      map[string]string
 	}{
 		{
 			name: "AddLabelsToEmptyMap",
-			uut:  &Lease{},
 			labels: map[string]string{
 				"containerd.io/gc.root": "2015-12-04T00:00:00Z",
 			},
@@ -42,10 +41,8 @@ func TestWithLabels(t *testing.T) {
 		},
 		{
 			name: "AddLabelsToNonEmptyMap",
-			uut: &Lease{
-				Labels: map[string]string{
-					"containerd.io/gc.expire": "2015-12-05T00:00:00Z",
-				},
+			initialLabels: 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",
@@ -62,9 +59,32 @@ func TestWithLabels(t *testing.T) {
 	for _, tc := range testcases {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			err := WithLabels(tc.labels)(tc.uut)
+			lease := newLease(tc.initialLabels)
+			err := WithLabels(tc.labels)(lease)
 			require.NoError(t, err)
-			assert.Equal(t, tc.uut.Labels, tc.expected)
+			assert.Equal(t, lease.Labels, tc.expected)
+		})
+	}
+	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
+}