refactor tests

Azure · Jan 13, 2025 · 78614ae · 78614ae
1 parent 94d1dac
commit 78614ae
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 130 deletions.
diff --git a/sdk/azidentity/default_azure_credential_test.go b/sdk/azidentity/default_azure_credential_test.go
@@ -10,7 +10,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -322,7 +321,6 @@ func TestDefaultAzureCredential_IMDS(t *testing.T) {
 						probed = true
 						require.Empty(t, hdr, "probe request shouldn't have Metadata header")
 						return &http.Response{
-							Body:       io.NopCloser(strings.NewReader("{}")),
 							StatusCode: http.StatusInternalServerError,
 						}
 					},
@@ -396,6 +394,84 @@ func TestDefaultAzureCredential_IMDS(t *testing.T) {
 	})
 }
 
+func TestDefaultAzureCredential_UnexpectedIMDSResponse(t *testing.T) {
+	before := defaultAzTokenProvider
+	defer func() { defaultAzTokenProvider = before }()
+	defaultAzTokenProvider = mockAzTokenProviderSuccess
+
+	const dockerDesktopPrefix = "connecting to 169.254.169.254:80: connecting to 169.254.169.254:80: dial tcp 169.254.169.254:80: connectex: A socket operation was attempted to an unreachable "
+	for _, test := range []struct {
+		desc string
+		res  [][]mock.ResponseOption
+	}{
+		{
+			"Docker Desktop",
+			[][]mock.ResponseOption{
+				{
+					mock.WithBody([]byte(dockerDesktopPrefix + "host.")),
+					mock.WithStatusCode(http.StatusForbidden),
+				},
+				{
+					mock.WithBody([]byte(dockerDesktopPrefix + "host.")),
+					mock.WithStatusCode(http.StatusForbidden),
+				},
+			},
+		},
+		{
+			"Docker Desktop",
+			[][]mock.ResponseOption{
+				{
+					mock.WithBody([]byte(dockerDesktopPrefix + "network.")),
+					mock.WithStatusCode(http.StatusForbidden),
+				},
+				{
+					mock.WithBody([]byte(dockerDesktopPrefix + "network.")),
+					mock.WithStatusCode(http.StatusForbidden),
+				},
+			},
+		},
+		{
+			"IMDS: no identity assigned",
+			[][]mock.ResponseOption{
+				{mock.WithStatusCode(http.StatusBadRequest)},
+				{
+					mock.WithBody([]byte(`{"error":"invalid_request","error_description":"Identity not found"}`)),
+					mock.WithStatusCode(http.StatusBadRequest),
+				},
+			},
+		},
+		{
+			"no token in response",
+			[][]mock.ResponseOption{
+				{mock.WithStatusCode(http.StatusOK)},
+				{mock.WithBody([]byte(`{"error": "no token here"}`)), mock.WithStatusCode(http.StatusOK)},
+			},
+		},
+		{
+			"non-JSON token response",
+			[][]mock.ResponseOption{
+				{mock.WithStatusCode(http.StatusOK)},
+				{mock.WithBody([]byte("not json")), mock.WithStatusCode(http.StatusOK)},
+			},
+		},
+	} {
+		t.Run(test.desc, func(t *testing.T) {
+			srv, close := mock.NewServer(mock.WithTransformAllRequestsToTestServerUrl())
+			defer close()
+			for _, res := range test.res {
+				srv.AppendResponse(res...)
+			}
+			c, err := NewDefaultAzureCredential(&DefaultAzureCredentialOptions{
+				ClientOptions: policy.ClientOptions{Transport: srv},
+			})
+			require.NoError(t, err)
+			tk, err := c.GetToken(ctx, testTRO)
+			require.NoError(t, err, "expected a token from AzureCLICredential")
+			require.Equal(t, tokenValue, tk.Token, "expected a token from AzureCLICredential")
+		})
+	}
+}
+
 func TestDefaultAzureCredential_UnsupportedMIClientID(t *testing.T) {
 	fail := true
 	before := defaultAzTokenProvider

diff --git a/sdk/azidentity/managed_identity_client_test.go b/sdk/azidentity/managed_identity_client_test.go
@@ -8,17 +8,13 @@ package azidentity
 
 import (
 	"context"
-	"errors"
-	"fmt"
 	"net/http"
 	"net/url"
 	"strings"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
-	"github.com/Azure/azure-sdk-for-go/sdk/internal/mock"
-	"github.com/stretchr/testify/require"
 )
 
 type userAgentValidatingPolicy struct {
@@ -77,127 +73,3 @@ func TestManagedIdentityClient_ApplicationID(t *testing.T) {
 		t.Fatal(err)
 	}
 }
-
-func TestManagedIdentityClient_IMDSErrors(t *testing.T) {
-	for _, test := range []struct {
-		body, desc string
-		code       int
-	}{
-		{
-			desc: "No identity assigned",
-			code: http.StatusBadRequest,
-			body: `{"error":"invalid_request","error_description":"Identity not found"}`,
-		},
-		{
-			desc: "Docker Desktop",
-			code: http.StatusForbidden,
-			body: "connecting to 169.254.169.254:80: connecting to 169.254.169.254:80: dial tcp 169.254.169.254:80: connectex: A socket operation was attempted to an unreachable network.",
-		},
-		{
-			desc: "Docker Desktop",
-			code: http.StatusForbidden,
-			body: "connecting to 169.254.169.254:80: connecting to 169.254.169.254:80: dial tcp 169.254.169.254:80: connectex: A socket operation was attempted to an unreachable host.",
-		},
-		{
-			desc: "Azure Container Instances",
-			code: http.StatusBadRequest,
-			body: "Required metadata header not specified or not correct",
-		},
-	} {
-		t.Run(fmt.Sprint(test.code), func(t *testing.T) {
-			srv, close := mock.NewServer(mock.WithTransformAllRequestsToTestServerUrl())
-			defer close()
-			srv.SetResponse(mock.WithBody([]byte(test.body)), mock.WithStatusCode(test.code))
-			client, err := newManagedIdentityClient(&ManagedIdentityCredentialOptions{
-				ClientOptions: azcore.ClientOptions{Transport: srv},
-			})
-			if err != nil {
-				t.Fatal(err)
-			}
-			_, err = client.authenticate(context.Background(), nil, testTRO.Scopes)
-			if err == nil {
-				t.Fatal("expected an error")
-			}
-			if actual := err.Error(); !strings.Contains(actual, test.body) {
-				t.Fatalf("expected response body in error, got %q", actual)
-			}
-			var unavailableErr credentialUnavailable
-			if !errors.As(err, &unavailableErr) {
-				t.Fatalf("expected %T, got %T", unavailableErr, err)
-			}
-		})
-	}
-}
-
-func TestManagedIdentityClient_IMDSProbeReturnsUnavailableError(t *testing.T) {
-	for _, test := range []struct {
-		desc string
-		res  [][]mock.ResponseOption
-	}{
-		{
-			"Azure Container Instance",
-			[][]mock.ResponseOption{
-				{
-					mock.WithBody([]byte("Required metadata header not specified or not correct")),
-					mock.WithStatusCode(http.StatusBadRequest),
-				},
-				{mock.WithBody([]byte("error")), mock.WithStatusCode(http.StatusBadRequest)},
-			},
-		},
-		{
-			"404",
-			[][]mock.ResponseOption{
-				{mock.WithStatusCode(http.StatusNotFound)},
-				{mock.WithStatusCode(http.StatusNotFound)},
-			},
-		},
-		{
-			"non-JSON token response",
-			[][]mock.ResponseOption{
-				{mock.WithBody([]byte("Required metadata header not specified or not correct"))},
-				{mock.WithBody([]byte("not json"))},
-			},
-		},
-		{
-			"no token in response",
-			[][]mock.ResponseOption{
-				{mock.WithBody([]byte("Required metadata header not specified or not correct"))},
-				{mock.WithBody([]byte(`{"error": "no token here"}`)), mock.WithStatusCode(http.StatusBadRequest)},
-			},
-		},
-	} {
-		t.Run(test.desc, func(t *testing.T) {
-			srv, close := mock.NewServer(mock.WithTransformAllRequestsToTestServerUrl())
-			defer close()
-			for _, res := range test.res {
-				srv.AppendResponse(res...)
-			}
-			c, err := newManagedIdentityClient(&ManagedIdentityCredentialOptions{
-				ClientOptions: azcore.ClientOptions{
-					// disabling retries simplifies the 404 test (404 is retriable for IMDS)
-					Retry:     policy.RetryOptions{MaxRetries: -1},
-					Transport: srv,
-				},
-				dac: true,
-			})
-			require.NoError(t, err)
-			_, err = c.authenticate(ctx, nil, testTRO.Scopes)
-			var cu credentialUnavailable
-			require.ErrorAs(t, err, &cu)
-
-			srv.AppendResponse(
-				mock.WithBody(accessTokenRespSuccess),
-				mock.WithPredicate(func(r *http.Request) bool {
-					if _, ok := r.Header["Metadata"]; !ok {
-						t.Error("client shouldn't send another probe after receiving a response")
-					}
-					return true
-				}),
-			)
-			srv.AppendResponse()
-			tk, err := c.authenticate(ctx, nil, testTRO.Scopes)
-			require.NoError(t, err)
-			require.Equal(t, tokenValue, tk.Token)
-		})
-	}
-}