@@ -4,7 +4,6 @@ import (

 	"net/http"

 	"net/http/httptest"

 	"runtime"

-	"strings"

 	"testing"

 	"github.com/docker/docker/api/server/httputils"

@@ -12,37 +11,16 @@ import (

 	"golang.org/x/net/context"

 )

-func TestVersionMiddleware(t *testing.T) {

-	handler := func(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {

-		if httputils.VersionFromContext(ctx) == "" {

-			t.Fatal("Expected version, got empty string")

-		}

-		return nil

-	}

-

+func TestVersionMiddlewareVersion(t *testing.T) {

 	defaultVersion := "1.10.0"

 	minVersion := "1.2.0"

-	m := NewVersionMiddleware(defaultVersion, defaultVersion, minVersion)

-	h := m.WrapHandler(handler)

-

-	req, _ := http.NewRequest("GET", "/containers/json", nil)

-	resp := httptest.NewRecorder()

-	ctx := context.Background()

-	if err := h(ctx, resp, req, map[string]string{}); err != nil {

-		t.Fatal(err)

-	}

-}

-

-func TestVersionMiddlewareVersionTooOld(t *testing.T) {

+	expectedVersion := defaultVersion

 	handler := func(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {

-		if httputils.VersionFromContext(ctx) == "" {

-			t.Fatal("Expected version, got empty string")

-		}

+		v := httputils.VersionFromContext(ctx)

+		assert.Equal(t, expectedVersion, v)

 		return nil

 	}

-	defaultVersion := "1.10.0"

-	minVersion := "1.2.0"

 	m := NewVersionMiddleware(defaultVersion, defaultVersion, minVersion)

 	h := m.WrapHandler(handler)

@@ -50,44 +28,45 @@ func TestVersionMiddlewareVersionTooOld(t *testing.T) {

 	resp := httptest.NewRecorder()

 	ctx := context.Background()

-	vars := map[string]string{"version": "0.1"}

-	err := h(ctx, resp, req, vars)

-

-	if !strings.Contains(err.Error(), "client version 0.1 is too old. Minimum supported API version is 1.2.0") {

-		t.Fatalf("Expected too old client error, got %v", err)

+	tests := []struct {

+		reqVersion      string

+		expectedVersion string

+		errString       string

+	}{

+		{

+			expectedVersion: "1.10.0",

+		},

+		{

+			reqVersion:      "1.9.0",

+			expectedVersion: "1.9.0",

+		},

+		{

+			reqVersion: "0.1",

+			errString:  "client version 0.1 is too old. Minimum supported API version is 1.2.0, please upgrade your client to a newer version",

+		},

+		{

+			reqVersion: "9999.9999",

+			errString:  "client version 9999.9999 is too new. Maximum supported API version is 1.10.0",

+		},

 	}

-}

-func TestVersionMiddlewareVersionTooNew(t *testing.T) {

-	handler := func(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {

-		if httputils.VersionFromContext(ctx) == "" {

-			t.Fatal("Expected version, got empty string")

-		}

-		return nil

-	}

+	for _, test := range tests {

+		expectedVersion = test.expectedVersion

-	defaultVersion := "1.10.0"

-	minVersion := "1.2.0"

-	m := NewVersionMiddleware(defaultVersion, defaultVersion, minVersion)

-	h := m.WrapHandler(handler)

+		err := h(ctx, resp, req, map[string]string{"version": test.reqVersion})

-	req, _ := http.NewRequest("GET", "/containers/json", nil)

-	resp := httptest.NewRecorder()

-	ctx := context.Background()

-

-	vars := map[string]string{"version": "9999.9999"}

-	err := h(ctx, resp, req, vars)

-

-	if !strings.Contains(err.Error(), "client version 9999.9999 is too new. Maximum supported API version is 1.10.0") {

-		t.Fatalf("Expected too new client error, got %v", err)

+		if test.errString != "" {

+			assert.EqualError(t, err, test.errString)

+		} else {

+			assert.NoError(t, err)

+		}

 	}

 }

 func TestVersionMiddlewareWithErrorsReturnsHeaders(t *testing.T) {

 	handler := func(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {

-		if httputils.VersionFromContext(ctx) == "" {

-			t.Fatal("Expected version, got empty string")

-		}

+		v := httputils.VersionFromContext(ctx)

+		assert.NotEmpty(t, v)

 		return nil

 	}

@@ -102,8 +81,8 @@ func TestVersionMiddlewareWithErrorsReturnsHeaders(t *testing.T) {

 	vars := map[string]string{"version": "0.1"}

 	err := h(ctx, resp, req, vars)

-

 	assert.Error(t, err)

+

 	hdr := resp.Result().Header

 	assert.Contains(t, hdr.Get("Server"), "Docker/"+defaultVersion)

 	assert.Contains(t, hdr.Get("Server"), runtime.GOOS)