Repository: Azure/azure-sdk-for-go
Pull request #17755 (2022-04-27 18:22:00)

Sync eng/common directory with azure-sdk-tools for PR 3238
Sync eng/common directory with azure-sdk-tools for PR https://github.com/Azure/azure-sdk-tools/pull/3238  See [eng/common workflow](https://github.com/Azure/azure-sdk-tools/blob/main/eng/common/README.md#workflow)


Patch:
diff --git a/sdk/data/azcosmos/CHANGELOG.md b/sdk/data/azcosmos/CHANGELOG.md
index 6bd050a4d5..e062a92998 100644
--- a/sdk/data/azcosmos/CHANGELOG.md
+++ b/sdk/data/azcosmos/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 ### Features Added
 * Added single partition query support.
+* Added Azure AD authentication support through `azcosmos.NewClient`
 
 ### Breaking Changes
 * This module now requires Go 1.18
diff --git a/sdk/data/azcosmos/README.md b/sdk/data/azcosmos/README.md
index ac0bd012e2..50792db7b4 100644
--- a/sdk/data/azcosmos/README.md
+++ b/sdk/data/azcosmos/README.md
@@ -44,13 +44,29 @@ The following section provides several code snippets covering some of the most c
 
 ### Create Cosmos Client
 
+The clients support different forms of authentication. The azcosmos library supports authorization via Azure Active Directory or an account key.
+
+**Using Azure Active Directory**
+
+```go
+import "github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+
+cred, err := azidentity.NewDefaultAzureCredential(nil)
+handle(err)
+client, err := azcosmos.NewClient("myAccountEndpointURL", cred, nil)
+handle(err)
+```
+
+**Using account keys**
+
 ```go
 const (
     cosmosDbEndpoint = "someEndpoint"
     cosmosDbKey = "someKey"
 )
 
-cred, _ := azcosmos.NewKeyCredential(cosmosDbKey)
+cred, err := azcosmos.NewKeyCredential(cosmosDbKey)
+handle(err)
 client, err := azcosmos.NewClientWithKey(cosmosDbEndpoint, cred, nil)
 handle(err)
 ```
diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go
index c9df942bf2..b1d62a73cb 100644
--- a/sdk/data/azcosmos/cosmos_client.go
+++ b/sdk/data/azcosmos/cosmos_client.go
@@ -7,9 +7,12 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"net/http"
+	"net/url"
 	"time"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming"
@@ -26,15 +29,27 @@ func (c *Client) Endpoint() string {
 	return c.endpoint
 }
 
-// NewClientWithKey creates a new instance of Cosmos client with the specified values. It uses the default pipeline configuration.
+// NewClientWithKey creates a new instance of Cosmos client with shared key authentication. It uses the default pipeline configuration.
 // endpoint - The cosmos service endpoint to use.
 // cred - The credential used to authenticate with the cosmos service.
 // options - Optional Cosmos client options.  Pass nil to accept default values.
 func NewClientWithKey(endpoint string, cred KeyCredential, o *ClientOptions) (*Client, error) {
-	return &Client{endpoint: endpoint, pipeline: newPipeline(cred, o)}, nil
+	return &Client{endpoint: endpoint, pipeline: newPipeline([]policy.Policy{newSharedKeyCredPolicy(cred)}, o)}, nil
 }
 
-func newPipeline(cred KeyCredential, options *ClientOptions) azruntime.Pipeline {
+// NewClient creates a new instance of Cosmos client with Azure AD access token authentication. It uses the default pipeline configuration.
+// endpoint - The cosmos service endpoint to use.
+// cred - The credential used to authenticate with the cosmos service.
+// options - Optional Cosmos client options.  Pass nil to accept default values.
+func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) (*Client, error) {
+	scope, err := createScopeFromEndpoint(endpoint)
+	if err != nil {
+		return nil, err
+	}
+	return &Client{endpoint: endpoint, pipeline: newPipeline([]policy.Policy{azruntime.NewBearerTokenPolicy(cred, scope, nil), &cosmosBearerTokenPolicy{}}, o)}, nil
+}
+
+func newPipeline(authPolicy []policy.Policy, options *ClientOptions) azruntime.Pipeline {
 	if options == nil {
 		options = &ClientOptions{}
 	}
@@ -42,15 +57,24 @@ func newPipeline(cred KeyCredential, options *ClientOptions) azruntime.Pipeline
 	return azruntime.NewPipeline("azcosmos", serviceLibVersion,
 		azruntime.PipelineOptions{
 			PerCall: []policy.Policy{
-				newSharedKeyCredPolicy(cred),
 				&headerPolicies{
 					enableContentResponseOnWrite: options.EnableContentResponseOnWrite,
 				},
 			},
+			PerRetry: authPolicy,
 		},
 		&options.ClientOptions)
 }
 
+func createScopeFromEndpoint(endpoint string) ([]string, error) {
+	u, err := url.Parse(endpoint)
+	if err != nil {
+		return nil, err
+	}
+
+	return []string{fmt.Sprintf("%s://%s/.default", u.Scheme, u.Hostname())}, nil
+}
+
 // NewDatabase returns a struct that represents a database and allows database level operations.
 // id - The id of the database.
 func (c *Client) NewDatabase(id string) (*DatabaseClient, error) {
diff --git a/sdk/data/azcosmos/cosmos_policy_bearer_token.go b/sdk/data/azcosmos/cosmos_policy_bearer_token.go
new file mode 100644
index 0000000000..dfba0e9b3c
--- /dev/null
+++ b/sdk/data/azcosmos/cosmos_policy_bearer_token.go
@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+)
+
+const lenBearerTokenPrefix = len("Bearer ")
+
+type cosmosBearerTokenPolicy struct {
+}
+
+func (b *cosmosBearerTokenPolicy) Do(req *policy.Request) (*http.Response, error) {
+	currentAuthorization := req.Raw().Header.Get(headerAuthorization)
+	if currentAuthorization == "" {
+		return nil, errors.New("authorization header is missing")
+	}
+
+	token := currentAuthorization[lenBearerTokenPrefix:]
+	req.Raw().Header.Set(headerAuthorization, fmt.Sprintf("type=aad&ver=1.0&sig=%v", token))
+	return req.Next()
+}
diff --git a/sdk/data/azcosmos/doc.go b/sdk/data/azcosmos/doc.go
index 30ccb34271..b8689fa183 100644
--- a/sdk/data/azcosmos/doc.go
+++ b/sdk/data/azcosmos/doc.go
@@ -12,6 +12,19 @@ The azcosmos package is capable of:
 
 Creating the Client
 
+Types of Credentials
+The clients support different forms of authentication. The azcosmos library supports
+authorization via Azure Active Directory or an account key.
+
+Using Azure Active Directory
+To create a client, you can use any of the TokenCredential implementations provided by `azidentity`.
+	cred, err := azidentity.NewClientSecretCredential("tenantId", "clientId", "clientSecret")
+	handle(err)
+	client, err := azcosmos.NewClient("myAccountEndpointURL", cred, nil)
+	handle(err)
+
+
+Using account keys
 To create a client, you will need the account's endpoint URL and a key credential.
 
 	cred, err := azcosmos.NewKeyCredential("myAccountKey")
diff --git a/sdk/data/azcosmos/go.mod b/sdk/data/azcosmos/go.mod
index 53a61a5260..f1f166c940 100644
--- a/sdk/data/azcosmos/go.mod
+++ b/sdk/data/azcosmos/go.mod
@@ -5,14 +5,22 @@ go 1.18
 require (
 	github.com/Azure/azure-sdk-for-go v63.2.0+incompatible
 	github.com/Azure/azure-sdk-for-go/sdk/azcore v0.23.0
+	github.com/Azure/azure-sdk-for-go/sdk/azidentity v0.14.0
 	github.com/Azure/azure-sdk-for-go/sdk/internal v0.9.2
 	github.com/stretchr/testify v1.7.0
 )
 
 require (
+	github.com/AzureAD/microsoft-authentication-library-for-go v0.4.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/golang-jwt/jwt v3.2.1+incompatible // indirect
+	github.com/google/uuid v1.1.1 // indirect
+	github.com/kylelemons/godebug v1.1.0 // indirect
+	github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897 // indirect
 	golang.org/x/net v0.0.0-20211015210444-4f30a5c0130f // indirect
+	golang.org/x/sys v0.0.0-20211019181941-9d821ace8654 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 )
diff --git a/sdk/data/azcosmos/go.sum b/sdk/data/azcosmos/go.sum
index 11b87f3d51..edd062b165 100644
--- a/sdk/data/azcosmos/go.sum
+++ b/sdk/data/azcosmos/go.sum
@@ -2,13 +2,28 @@ github.com/Azure/azure-sdk-for-go v63.2.0+incompatible h1:OIqkK/zTGqVUuzpEvY0B1Y
 github.com/Azure/azure-sdk-for-go v63.2.0+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v0.23.0 h1:D7l5jspkc4kwBYRWoZE4DQnu6LVpLwDsMZjBKS4wZLQ=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v0.23.0/go.mod h1:w5pDIZuawUmY3Bj4tVx3Xb8KS96ToB0j315w9rqpAg0=
+github.com/Azure/azure-sdk-for-go/sdk/azidentity v0.14.0 h1:NVS/4LOQfkBpk+B1VopIzv1ptmYeEskA8w/3K/w7vjo=
+github.com/Azure/azure-sdk-for-go/sdk/azidentity v0.14.0/go.mod h1:RG0cZndeZM17StwohYclmcXSr4oOJ8b1I5hB8llIc6Y=
 github.com/Azure/azure-sdk-for-go/sdk/internal v0.9.2 h1:Px2KVERcYEg2Lv25AqC2hVr0xUWaq94wuEObLIkYzmA=
 github.com/Azure/azure-sdk-for-go/sdk/internal v0.9.2/go.mod h1:CdSJQNNzZhCkwDaV27XV1w48ZBPtxe7mlrZAsPNxD5g=
+github.com/AzureAD/microsoft-authentication-library-for-go v0.4.0 h1:WVsrXCnHlDDX8ls+tootqRE87/hL9S/g4ewig9RsD/c=
+github.com/AzureAD/microsoft-authentication-library-for-go v0.4.0/go.mod h1:Vt9sXTKwMyGcOxSmLDMnGPgqsUg7m8pe215qMLrDXw4=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dnaeon/go-vcr v1.1.0 h1:ReYa/UBrRyQdant9B4fNHGoCNKw6qh6P0fsdGmZpR7c=
 github.com/dnaeon/go-vcr v1.1.0/go.mod h1:M7tiix8f0r6mKKJ3Yq/kqU1OYf3MnfmBWVbPx/yU9ko=
+github.com/golang-jwt/jwt v3.2.1+incompatible h1:73Z+4BJcrTC+KczS6WvTPvRGOp1WmfEP4Q1lOd9Z/+c=
+github.com/golang-jwt/jwt v3.2.1+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I=
+github.com/golang-jwt/jwt/v4 v4.2.0 h1:besgBTC8w8HjP6NzQdxwKH9Z5oQMZ24ThTrHp3cZ8eU=
+github.com/google/uuid v1.1.1 h1:Gkbcsh/GbpXz7lPftLA3P6TYMwjCLYm83jiFQZF/3gY=
+github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
+github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
 github.com/modocache/gover v0.0.0-20171022184752-b58185e213c5/go.mod h1:caMODM3PzxT8aQXRPkAt8xlV/e7d7w8GM5g0fa5F0D8=
+github.com/montanaflynn/stats v0.6.6/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow=
+github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4 h1:Qj1ukM4GlMWXNdMBuXcXfz/Kw9s1qm0CLY32QxuSImI=
+github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4/go.mod h1:N6UoU20jOqggOuDwUaBQpluzLNDqif3kq9z2wpdYEfQ=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -16,6 +31,8 @@ github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5Cc
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
+golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897 h1:pLI5jrR7OSLijeIDcmRxNmw2api+jEfxLoykJVice/E=
+golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20201010224723-4f7140c49acb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20211015210444-4f30a5c0130f h1:OfiFi4JbukWwe3lzw+xunroH1mnC1e2Gy5cxNJApiSY=
@@ -23,6 +40,9 @@ golang.org/x/net v0.0.0-20211015210444-4f30a5c0130f/go.mod h1:9nx3DQGgdP8bBQD5qx
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20211019181941-9d821ace8654 h1:id054HUawV2/6IGm2IV8KZQjqtwAOo2CYlOToYqa0d0=
+golang.org/x/sys v0.0.0-20211019181941-9d821ace8654/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
@@ -31,6 +51,7 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b h1:h8qDotaEPuJATrMmW04NCwg7v22aHH28wwpauUhK9Oo=


Test patch:
diff --git a/eng/common/testproxy/docker-start-proxy.ps1 b/eng/common/testproxy/docker-start-proxy.ps1
index c34c9d38d8..5cdeebf590 100644
--- a/eng/common/testproxy/docker-start-proxy.ps1
+++ b/eng/common/testproxy/docker-start-proxy.ps1
@@ -38,7 +38,7 @@ catch {
 
 $SELECTED_IMAGE_TAG = $(Get-Content "$PSScriptRoot/target_version.txt" -Raw).Trim()
 $CONTAINER_NAME = "ambitious_azsdk_test_proxy"
-$IMAGE_SOURCE = "azsdkengsys.azurecr.io/engsys/testproxy:${SELECTED_IMAGE_TAG}"
+$IMAGE_SOURCE = "azsdkengsys.azurecr.io/engsys/test-proxy:${SELECTED_IMAGE_TAG}"
 
 if($VersionOverride) { 
     Write-Host "Overriding default target proxy version of '$SELECTED_IMAGE_TAG' with override $VersionOverride."
diff --git a/eng/common/testproxy/target_version.txt b/eng/common/testproxy/target_version.txt
index a72b151d94..86b34733bf 100644
--- a/eng/common/testproxy/target_version.txt
+++ b/eng/common/testproxy/target_version.txt
@@ -1,1 +1,1 @@
-1.0.0-dev.20220426.1
+1.0.0-dev.20220427.1
diff --git a/eng/pipelines/templates/steps/build-test.yml b/eng/pipelines/templates/steps/build-test.yml
index b06ca963f8..17dc761ac5 100644
--- a/eng/pipelines/templates/steps/build-test.yml
+++ b/eng/pipelines/templates/steps/build-test.yml
@@ -48,9 +48,6 @@ steps:
 
   - ${{ if eq(parameters.TestProxy, true) }}:
     - template: /eng/common/testproxy/test-proxy-tool.yml
-      parameters:
-        targetVersion: '1.0.0-dev.20220224.2'
-
 
   - task: PowerShell@2
     displayName: 'Run Tests'
diff --git a/sdk/data/azcosmos/cosmos_client_test.go b/sdk/data/azcosmos/cosmos_client_test.go
index 1eab3dcba3..6bda337e82 100644
--- a/sdk/data/azcosmos/cosmos_client_test.go
+++ b/sdk/data/azcosmos/cosmos_client_test.go
@@ -354,6 +354,22 @@ func TestSendQuery(t *testing.T) {
 	}
 }
 
+func TestCreateScopeFromEndpoint(t *testing.T) {
+	url := "https://foo.documents.azure.com:443/"
+	scope, err := createScopeFromEndpoint(url)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if scope[0] != "https://foo.documents.azure.com/.default" {
+		t.Errorf("Expected %v, but got %v", "https://foo.documents.azure.com/.default", scope[0])
+	}
+
+	if len(scope) != 1 {
+		t.Errorf("Expected %v, but got %v", 1, len(scope))
+	}
+}
+
 type pipelineVerifier struct {
 	requests []pipelineVerifierRequest
 }
diff --git a/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go b/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go
new file mode 100644
index 0000000000..e5d3eecade
--- /dev/null
+++ b/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go
@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/mock"
+)
+
+func TestConvertBearerToken(t *testing.T) {
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.SetResponse(mock.WithStatusCode(http.StatusOK))
+
+	verifier := bearerTokenVerify{}
+	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{PerCall: []policy.Policy{&mockAuthPolicy{}, &cosmosBearerTokenPolicy{}, &verifier}}, &policy.ClientOptions{Transport: srv})
+	req, err := azruntime.NewRequest(context.Background(), http.MethodGet, srv.URL())
+	req.SetOperationValue(pipelineRequestOptions{
+		isWriteOperation: true,
+	})
+
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	_, err = pl.Do(req)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if verifier.authHeaderContent != "type=aad&ver=1.0&sig=this is a test token" {
+		t.Fatalf("Expected auth header content to be 'type=aad&ver=1.0&sig=this is a test token', got %s", verifier.authHeaderContent)
+	}
+}
+
+type bearerTokenVerify struct {
+	authHeaderContent string
+}
+
+func (p *bearerTokenVerify) Do(req *policy.Request) (*http.Response, error) {
+	p.authHeaderContent = req.Raw().Header.Get(headerAuthorization)
+
+	return req.Next()
+}
+
+type mockAuthPolicy struct{}
+
+func (p *mockAuthPolicy) Do(req *policy.Request) (*http.Response, error) {
+	req.Raw().Header.Set(headerAuthorization, "Bearer this is a test token")
+
+	return req.Next()
+}
diff --git a/sdk/data/azcosmos/emulator_cosmos_aad_test.go b/sdk/data/azcosmos/emulator_cosmos_aad_test.go
new file mode 100644
index 0000000000..8a1502507f
--- /dev/null
+++ b/sdk/data/azcosmos/emulator_cosmos_aad_test.go
@@ -0,0 +1,140 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+)
+
+func TestAAD(t *testing.T) {
+	emulatorTests := newEmulatorTests(t)
+	client := emulatorTests.getClient(t)
+
+	database := emulatorTests.createDatabase(t, context.TODO(), client, "aadTest")
+	defer emulatorTests.deleteDatabase(t, context.TODO(), database)
+	properties := ContainerProperties{
+		ID: "aContainer",
+		PartitionKeyDefinition: PartitionKeyDefinition{
+			Paths: []string{"/id"},
+		},
+	}
+
+	_, err := database.CreateContainer(context.TODO(), properties, nil)
+	if err != nil {
+		t.Fatalf("Failed to create container: %v", err)
+	}
+
+	aadClient := emulatorTests.getAadClient(t)
+
+	item := map[string]string{
+		"id":    "1",
+		"value": "2",
+	}
+
+	container, _ := aadClient.NewContainer("aadTest", "aContainer")
+	pk := NewPartitionKeyString("1")
+
+	marshalled, err := json.Marshal(item)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	itemResponse, err := container.CreateItem(context.TODO(), pk, marshalled, nil)
+	if err != nil {
+		t.Fatalf("Failed to create item: %v", err)
+	}
+
+	if itemResponse.SessionToken == "" {
+		t.Fatalf("Session token is empty")
+	}
+
+	// No content on write by default
+	if len(itemResponse.Value) != 0 {
+		t.Fatalf("Expected empty response, got %v", itemResponse.Value)
+	}
+
+	itemResponse, err = container.ReadItem(context.TODO(), pk, "1", nil)
+	if err != nil {
+		t.Fatalf("Failed to read item: %v", err)
+	}
+
+	if len(itemResponse.Value) == 0 {
+		t.Fatalf("Expected non-empty response, got %v", itemResponse.Value)
+	}
+
+	var itemResponseBody map[string]interface{}
+	err = json.Unmarshal(itemResponse.Value, &itemResponseBody)
+	if err != nil {
+		t.Fatalf("Failed to unmarshal item response: %v", err)
+	}
+	if itemResponseBody["id"] != "1" {
+		t.Fatalf("Expected id to be 1, got %v", itemResponseBody["id"])
+	}
+	if itemResponseBody["value"] != "2" {
+		t.Fatalf("Expected value to be 2, got %v", itemResponseBody["value"])
+	}
+
+	item["value"] = "3"
+	marshalled, err = json.Marshal(item)
+	if err != nil {
+		t.Fatal(err)
+	}
+	itemResponse, err = container.ReplaceItem(context.TODO(), pk, "1", marshalled, &ItemOptions{EnableContentResponseOnWrite: true})
+	if err != nil {
+		t.Fatalf("Failed to replace item: %v", err)
+	}
+
+	// Explicitly requesting body on write
+	if len(itemResponse.Value) == 0 {
+		t.Fatalf("Expected non-empty response, got %v", itemResponse.Value)
+	}
+
+	err = json.Unmarshal(itemResponse.Value, &itemResponseBody)
+	if err != nil {
+		t.Fatalf("Failed to unmarshal item response: %v", err)
+	}
+	if itemResponseBody["id"] != "1" {
+		t.Fatalf("Expected id to be 1, got %v", itemResponseBody["id"])
+	}
+	if itemResponseBody["value"] != "3" {
+		t.Fatalf("Expected value to be 3, got %v", itemResponseBody["value"])
+	}
+
+	item["value"] = "4"
+	marshalled, err = json.Marshal(item)
+	if err != nil {
+		t.Fatal(err)
+	}
+	itemResponse, err = container.UpsertItem(context.TODO(), pk, marshalled, &ItemOptions{EnableContentResponseOnWrite: true})
+	if err != nil {
+		t.Fatalf("Failed to upsert item: %v", err)
+	}
+
+	// Explicitly requesting body on write
+	if len(itemResponse.Value) == 0 {
+		t.Fatalf("Expected non-empty response, got %v", itemResponse.Value)
+	}
+
+	err = json.Unmarshal(itemResponse.Value, &itemResponseBody)
+	if err != nil {
+		t.Fatalf("Failed to unmarshal item response: %v", err)
+	}
+	if itemResponseBody["id"] != "1" {
+		t.Fatalf("Expected id to be 1, got %v", itemResponseBody["id"])
+	}
+	if itemResponseBody["value"] != "4" {
+		t.Fatalf("Expected value to be 4, got %v", itemResponseBody["value"])
+	}
+
+	itemResponse, err = container.DeleteItem(context.TODO(), pk, "1", nil)
+	if err != nil {
+		t.Fatalf("Failed to replace item: %v", err)
+	}
+
+	if len(itemResponse.Value) != 0 {
+		t.Fatalf("Expected empty response, got %v", itemResponse.Value)
+	}
+}
diff --git a/sdk/data/azcosmos/emulator_tests.go b/sdk/data/azcosmos/emulator_tests.go
index d0a1793995..c8994ab939 100644
--- a/sdk/data/azcosmos/emulator_tests.go
+++ b/sdk/data/azcosmos/emulator_tests.go
@@ -5,8 +5,14 @@ package azcosmos
 
 import (
 	"context"
+	"encoding/base64"
 	"os"
+	"strconv"
 	"testing"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 )
 
 type emulatorTests struct {
@@ -36,6 +42,16 @@ func (e *emulatorTests) getClient(t *testing.T) *Client {
 	return client
 }
 
+func (e *emulatorTests) getAadClient(t *testing.T) *Client {
+	cred := &emulatorTokenCredential{}
+	client, err := NewClient(e.host, cred, nil)
+	if err != nil {
+		t.Fatalf("Failed to create client: %v", err)
+	}
+
+	return client
+}
+
 func (e *emulatorTests) createDatabase(
 	t *testing.T,
 	ctx context.Context,
@@ -64,3 +80,47 @@ func (e *emulatorTests) deleteDatabase(
 		t.Fatalf("Failed to delete database: %v", err)
 	}
 }
+
+type emulatorTokenCredential struct {
+}
+
+func (c *emulatorTokenCredential) GetToken(ctx context.Context, options policy.TokenRequestOptions) (*azcore.AccessToken, error) {
+	header := `{"typ":"JWT","alg":"RS256","x5t":"CosmosEmulatorPrimaryMaster","kid":"CosmosEmulatorPrimaryMaster"}`
+	unixNow := time.Now().Unix()
+	expiration := unixNow + 7200
+	payload := `{ 
+		"appid":"localhost", 
+		"aio":"",
+		"appidacr":"1",
+		"idp": "https://localhost:8081/",
+		"oid": "96313034-4739-43cb-93cd-74193adbe5b6",
+		"rh": "",
+		"sub": "localhost",
+		"tid": "EmulatorFederation",
+		"uti": "",
+		"ver": "1.0",
+		"scp": "user_impersonation",
+		"groups":[ 
+			"7ce1d003-4cb3-4879-b7c5-74062a35c66e",
+			"e99ff30c-c229-4c67-ab29-30a6aebc3e58",
+			"5549bb62-c77b-4305-bda9-9ec66b85d9e4",
+			"c44fd685-5c58-452c-aaf7-13ce75184f65",
+			"be895215-eab5-43b7-9536-9ef8fe130330"], 
+		"nbf":` + strconv.FormatInt(unixNow, 10) + `, 
+		"exp":` + strconv.FormatInt(expiration, 10) + `, 
+		"iat":` + strconv.FormatInt(unixNow, 10) + `,
+		"iss":"https://sts.fake-issuer.net/7b1999a1-dfd7-440e-8204-00170979b984",
+		"aud":"https://localhost.localhost" 
+	}`
+
+	headerBase64 := base64.RawURLEncoding.EncodeToString([]byte(header))
+	payloadBase64 := base64.RawURLEncoding.EncodeToString([]byte(payload))
+	masterKeyBase64 := base64.RawURLEncoding.EncodeToString([]byte("C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="))
+
+	token := headerBase64 + "." + payloadBase64 + "." + masterKeyBase64
+
+	return &azcore.AccessToken{
+		Token:     token,
+		ExpiresOn: time.Unix(expiration, 0),
+	}, nil
+}
diff --git a/sdk/data/azcosmos/example_test.go b/sdk/data/azcosmos/example_test.go
index 9d927a8f9f..32d5161468 100644
--- a/sdk/data/azcosmos/example_test.go
+++ b/sdk/data/azcosmos/example_test.go
@@ -12,9 +12,32 @@ import (
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
 )
 
+func ExampleNewClient() {
+	endpoint, ok := os.LookupEnv("AZURE_COSMOS_ENDPOINT")
+	if !ok {
+		panic("AZURE_COSMOS_ENDPOINT could not be found")
+	}
+
+	// Obtain a TokenCredential for the current environment
+	// Alternatively, you could use any of the other credential types
+	// For example, azidentity.NewClientSecretCredential("tenantId", "clientId", "clientSecret")
+	cred, err := azidentity.NewDefaultAzureCredential(nil)
+	if err != nil {
+		panic(err)
+	}
+
+	client, err := azcosmos.NewClient(endpoint, cred, nil)
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println(client)
+}
+
 func ExampleNewClientWithKey() {
 	endpoint, ok := os.LookupEnv("AZURE_COSMOS_ENDPOINT")
 	if !ok {