Allow test runners to wrap contexts

Let the test runners choose the namespaces and
wrap the contexts. This allows the test suite to create
multiple contexts without worrying about namespacing
or leasing in the contexts.

Signed-off-by: Derek McGowan <derek@mcgstyle.net>

content/testsuite/testsuite.go

@@ -14,7 +14,6 @@ import (
 	"time"
 
 	"github.com/containerd/containerd/content"
-	"github.com/containerd/containerd/namespaces"
 	"github.com/containerd/containerd/testutil"
 	"github.com/gotestyourself/gotestyourself/assert"
 	digest "github.com/opencontainers/go-digest"
@@ -34,9 +33,33 @@ func ContentSuite(t *testing.T, name string, storeFn func(ctx context.Context, r
 	t.Run("Labels", makeTest(t, name, storeFn, checkLabels))
 }
 
+// ContextWrapper is used to decorate new context used inside the test
+// before using the context on the content store.
+// This can be used to support leasing and multiple namespaces tests.
+type ContextWrapper func(ctx context.Context) (context.Context, func() error, error)
+
+type wrapperKey struct{}
+
+// SetContextWrapper sets the wrapper on the context for deriving
+// new test contexts from the context.
+func SetContextWrapper(ctx context.Context, w ContextWrapper) context.Context {
+	return context.WithValue(ctx, wrapperKey{}, w)
+}
+
+type nameKey struct{}
+
+// Name gets the test name from the context
+func Name(ctx context.Context) string {
+	name, ok := ctx.Value(nameKey{}).(string)
+	if !ok {
+		return ""
+	}
+	return name
+}
+
 func makeTest(t *testing.T, name string, storeFn func(ctx context.Context, root string) (context.Context, content.Store, func() error, error), fn func(ctx context.Context, t *testing.T, cs content.Store)) func(t *testing.T) {
 	return func(t *testing.T) {
-		ctx := namespaces.WithNamespace(context.Background(), name)
+		ctx := context.WithValue(context.Background(), nameKey{}, name)
 
 		tmpDir, err := ioutil.TempDir("", "content-suite-"+name+"-")
 		if err != nil {
@@ -54,6 +77,20 @@ func makeTest(t *testing.T, name string, storeFn func(ctx context.Context, root
 			}
 		}()
 
+		w, ok := ctx.Value(wrapperKey{}).(ContextWrapper)
+		if ok {
+			var done func() error
+			ctx, done, err = w(ctx)
+			if err != nil {
+				t.Fatalf("Error wrapping context: %+v", err)
+			}
+			defer func() {
+				if err := done(); err != nil && !t.Failed() {
+					t.Fatalf("Wrapper release failed: %+v", err)
+				}
+			}()
+		}
+
 		defer testutil.DumpDirOnFailure(t, tmpDir)
 		fn(ctx, t, cs)
 	}

content_test.go

@@ -2,11 +2,14 @@ package containerd
 
 import (
 	"context"
+	"fmt"
+	"sync/atomic"
 	"testing"
 
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/content/testsuite"
 	"github.com/containerd/containerd/errdefs"
+	"github.com/containerd/containerd/namespaces"
 	"github.com/pkg/errors"
 )
 
@@ -15,13 +18,24 @@ func newContentStore(ctx context.Context, root string) (context.Context, content
 	if err != nil {
 		return nil, nil, nil, err
 	}
-	ctx, releaselease, err := client.WithLease(ctx)
-	if err != nil {
-		return nil, nil, nil, err
+
+	var (
+		count uint64
+		cs    = client.ContentStore()
+		name  = testsuite.Name(ctx)
+	)
+
+	wrap := func(ctx context.Context) (context.Context, func() error, error) {
+		n := atomic.AddUint64(&count, 1)
+		ctx = namespaces.WithNamespace(ctx, fmt.Sprintf("%s-n%d", name, n))
+		return client.WithLease(ctx)
 	}
-	cs := client.ContentStore()
+
+	ctx = testsuite.SetContextWrapper(ctx, wrap)
 
 	return ctx, cs, func() error {
+		for i := uint64(1); i <= count; i++ {
+			ctx = namespaces.WithNamespace(ctx, fmt.Sprintf("%s-n%d", name, i))
 			statuses, err := cs.ListStatuses(ctx)
 			if err != nil {
 				return err
@@ -31,8 +45,7 @@ func newContentStore(ctx context.Context, root string) (context.Context, content
 					return errors.Wrapf(err, "failed to abort %s", st.Ref)
 				}
 			}
-		releaselease()
-		return cs.Walk(ctx, func(info content.Info) error {
+			err = cs.Walk(ctx, func(info content.Info) error {
 				if err := cs.Delete(ctx, info.Digest); err != nil {
 					if errdefs.IsNotFound(err) {
 						return nil
@@ -42,6 +55,11 @@ func newContentStore(ctx context.Context, root string) (context.Context, content
 				}
 				return nil
 			})
+			if err != nil {
+				return err
+			}
+		}
+		return nil
 
 	}, nil
 }

metadata/content_test.go

@@ -3,7 +3,9 @@ package metadata
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"path/filepath"
+	"sync/atomic"
 	"testing"
 
 	"github.com/boltdb/bolt"
@@ -29,6 +31,16 @@ func createContentStore(ctx context.Context, root string) (context.Context, cont
 		return nil, nil, nil, err
 	}
 
+	var (
+		count uint64
+		name  = testsuite.Name(ctx)
+	)
+	wrap := func(ctx context.Context) (context.Context, func() error, error) {
+		n := atomic.AddUint64(&count, 1)
+		return namespaces.WithNamespace(ctx, fmt.Sprintf("%s-n%d", name, n)), func() error { return nil }, nil
+	}
+	ctx = testsuite.SetContextWrapper(ctx, wrap)
+
 	return ctx, NewDB(db, cs, nil).ContentStore(), func() error {
 		return db.Close()
 	}, nil

namespaces/context.go

@@ -1,11 +1,11 @@
 package namespaces
 
 import (
+	"context"
 	"os"
 
 	"github.com/containerd/containerd/errdefs"
 	"github.com/pkg/errors"
-	"golang.org/x/net/context"
 )
 
 const (