Add AAD integration tests for Cosmos DB

simorenoh · Copilot · simorenoh · commit 67450e3a084b · 2026-06-30T10:01:11.000-04:00
Adds the first AAD-authenticated integration test coverage for the Cosmos SDK, closing a GA-readiness gap where no test exercised the AAD token path.

- Add CosmosEmulatorCredential, a TokenCredential that mints the emulator's fake JWT, plus a CredentialRecorder so tests can assert the AAD path (not key auth) was used and which scope was requested.
- Add build_aad_client_from_env / TestRunContext::aad_client(), selecting the credential by host: emulator uses CosmosEmulatorCredential; live accounts use azure_core_test credentials from env (AzurePipelinesCredential in CI), matching the principal the test bicep grants the data-plane RBAC role to.
- Add cosmos_aad.rs with two dual-client tests: aad_item_crud_roundtrip (create/read/replace/query/delete via AAD) and aad_read_container_metadata. The key client handles db/container setup and cleanup since data-plane RBAC cannot create them.
- Enable AAD on the Windows emulator launch in Invoke-CosmosTestSetup.ps1 so the gated tests pass on the Windows Build leg (the Docker path already enabled it).

The tests are gated on test_category=emulator, so they run on both the Build-stage emulator and the live SingleWrite legs, giving real-AAD coverage on scheduled/weekly live runs at no extra config.

Co-authored-by: Copilot &lt;223556219+Copilot@users.noreply.github.com&gt;
diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_aad.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_aad.rs
@@ -0,0 +1,189 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+//! Integration tests exercising Entra ID (AAD) authentication against Azure
+//! Cosmos DB.
+//!
+//! These tests use a **dual-client** pattern: a key-auth client (provided by the
+//! framework) performs database/container management, while a separate
+//! AAD-authenticated client performs all data-plane item operations. This
+//! mirrors the data-plane RBAC role provisioned in `test-resources.bicep`, which
+//! grants item/metadata data actions but **not** management-plane permissions.
+//!
+//! Because the standard `test_category="emulator"` gate is used, the same tests
+//! run against the local emulator (Build stage, with the emulator started using
+//! `/enableaadauthentication`) and against live accounts (LiveTest stage), where
+//! the framework selects a real Entra ID credential via
+//! `azure_core_test::credentials::from_env`.
+
+use super::framework;
+
+use azure_core::http::StatusCode;
+use azure_core::Uuid;
+use azure_data_cosmos::models::ContainerProperties;
+use azure_data_cosmos::query::FeedScope;
+use azure_data_cosmos::{PartitionKey, Query};
+use framework::{TestClient, TestRunContext};
+use futures::TryStreamExt;
+use serde::{Deserialize, Serialize};
+use std::error::Error;
+
+/// The scope the Cosmos driver requests when acquiring an AAD token.
+const COSMOS_AAD_SCOPE: &str = "https://cosmos.azure.com/.default";
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
+struct AadTestItem {
+    id: String,
+    partition_key: String,
+    value: i64,
+}
+
+/// Drives a full item CRUD round-trip through an AAD-authenticated client.
+///
+/// Setup (database + container) and teardown run through the framework's
+/// key-auth client; only the item operations use the AAD client. On the
+/// emulator we additionally assert the bespoke fake-JWT credential was actually
+/// invoked for the Cosmos scope, guarding against silently exercising key auth.
+#[tokio::test]
+#[cfg_attr(
+    not(any(test_category = "emulator", test_category = "emulator_vnext")),
+    ignore = "requires test_category 'emulator' or 'emulator_vnext'"
+)]
+pub async fn aad_item_crud_roundtrip() -> Result<(), Box<dyn Error>> {
+    TestClient::run_with_unique_db(
+        async |run_context: &TestRunContext, db_client| {
+            // Key client creates the container (management-plane operation).
+            let container_id = format!("aad-container-{}", Uuid::new_v4());
+            run_context
+                .create_container(
+                    db_client,
+                    ContainerProperties::new(container_id.clone(), "/partition_key".into()),
+                    None,
+                )
+                .await?;
+
+            // Build the AAD-authenticated client and address the same container.
+            let (aad_client, recorder) = run_context.aad_client().await?;
+            let aad_container = aad_client
+                .database_client(db_client.id())
+                .container_client(&container_id)
+                .await?;
+
+            let unique = Uuid::new_v4().to_string();
+            let pk = format!("pk-{unique}");
+            let item_id = format!("item-{unique}");
+            let mut item = AadTestItem {
+                id: item_id.clone(),
+                partition_key: pk.clone(),
+                value: 1,
+            };
+
+            // Create via AAD.
+            let create = aad_container
+                .create_item(&pk, &item_id, &item, None)
+                .await?;
+            assert_eq!(
+                create.status(),
+                StatusCode::Created,
+                "AAD create_item should succeed"
+            );
+
+            // Point read via AAD.
+            let read = aad_container.read_item(&pk, &item_id, None).await?;
+            assert_eq!(read.status(), StatusCode::Ok);
+            assert_eq!(
+                read.into_model::<AadTestItem>()?,
+                item,
+                "round-tripped item should match"
+            );
+
+            // Replace via AAD.
+            item.value = 2;
+            let replace = aad_container
+                .replace_item(&pk, &item_id, &item, None)
+                .await?;
+            assert_eq!(replace.status(), StatusCode::Ok);
+
+            // Query via AAD, scoped to the item's partition.
+            let query =
+                Query::from("SELECT * FROM c WHERE c.id = @id").with_parameter("@id", &item_id)?;
+            let found: Vec<AadTestItem> = aad_container
+                .query_items::<AadTestItem>(
+                    query,
+                    FeedScope::partition(PartitionKey::from(&pk)),
+                    None,
+                )
+                .await?
+                .try_collect()
+                .await?;
+            assert_eq!(found.len(), 1, "query should return exactly the one item");
+            assert_eq!(found[0], item);
+
+            // Delete via AAD.
+            let delete = aad_container.delete_item(&pk, &item_id, None).await?;
+            assert_eq!(delete.status(), StatusCode::NoContent);
+
+            // On the emulator, prove the AAD credential was actually exercised.
+            if let Some(recorder) = recorder {
+                assert!(
+                    recorder.call_count() > 0,
+                    "emulator AAD credential get_token was never called"
+                );
+                assert!(
+                    recorder.requested_scope(COSMOS_AAD_SCOPE),
+                    "expected the Cosmos scope ({COSMOS_AAD_SCOPE}) to be requested, got {:?}",
+                    recorder.requested_scopes()
+                );
+            }
+
+            Ok(())
+        },
+        None,
+    )
+    .await
+}
+
+/// Verifies an AAD-authenticated client can read container metadata, exercising
+/// the `readMetadata` data action the SDK requires on its first request.
+#[tokio::test]
+#[cfg_attr(
+    not(any(test_category = "emulator", test_category = "emulator_vnext")),
+    ignore = "requires test_category 'emulator' or 'emulator_vnext'"
+)]
+pub async fn aad_read_container_metadata() -> Result<(), Box<dyn Error>> {
+    TestClient::run_with_unique_db(
+        async |run_context: &TestRunContext, db_client| {
+            let container_id = format!("aad-meta-{}", Uuid::new_v4());
+            run_context
+                .create_container(
+                    db_client,
+                    ContainerProperties::new(container_id.clone(), "/partition_key".into()),
+                    None,
+                )
+                .await?;
+
+            let (aad_client, recorder) = run_context.aad_client().await?;
+            let aad_container = aad_client
+                .database_client(db_client.id())
+                .container_client(&container_id)
+                .await?;
+
+            let properties = aad_container.read(None).await?.into_model()?;
+            assert_eq!(
+                properties.id, container_id,
+                "AAD container read should return the same container"
+            );
+
+            if let Some(recorder) = recorder {
+                assert!(
+                    recorder.requested_scope(COSMOS_AAD_SCOPE),
+                    "expected the Cosmos scope to be requested via AAD"
+                );
+            }
+
+            Ok(())
+        },
+        None,
+    )
+    .await
+}
diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/mod.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/mod.rs
@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+mod cosmos_aad;
 mod cosmos_backup_endpoints;
 mod cosmos_batch;
 mod cosmos_containers;
diff --git a/sdk/cosmos/azure_data_cosmos/tests/framework/emulator_credential.rs b/sdk/cosmos/azure_data_cosmos/tests/framework/emulator_credential.rs
@@ -0,0 +1,158 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+//! A [`TokenCredential`] that mints the fake JWT the Azure Cosmos DB emulator
+//! accepts when started with `/enableaadauthentication`.
+//!
+//! The emulator does not validate tokens against Entra ID. Instead it accepts a
+//! self-signed JWT whose signature segment is the emulator's well-known master
+//! key (encoded as base64url, no padding). This mirrors the
+//! `CosmosEmulatorCredential` used by the Python SDK's AAD emulator tests.
+//!
+//! This credential is for **tests only**. Never use it against a real account.
+
+use azure_core::credentials::{AccessToken, TokenCredential, TokenRequestOptions};
+use azure_core::time::OffsetDateTime;
+use base64::engine::general_purpose::URL_SAFE_NO_PAD;
+use base64::Engine as _;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+
+/// The well-known master key the Cosmos DB emulator ships with. The emulator
+/// validates the JWT signature segment against this exact string.
+pub const EMULATOR_MASTER_KEY: &str =
+    "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==";
+
+/// Records how a [`CosmosEmulatorCredential`] was exercised so tests can assert
+/// the AAD path (and not key auth) was actually used.
+#[derive(Clone, Default)]
+pub struct CredentialRecorder {
+    call_count: Arc<AtomicUsize>,
+    scopes: Arc<Mutex<Vec<String>>>,
+}
+
+impl CredentialRecorder {
+    /// Number of times `get_token` was called.
+    pub fn call_count(&self) -> usize {
+        self.call_count.load(Ordering::SeqCst)
+    }
+
+    /// All scopes that were requested across every `get_token` call.
+    pub fn requested_scopes(&self) -> Vec<String> {
+        self.scopes.lock().unwrap().clone()
+    }
+
+    /// Returns `true` if any `get_token` call requested `scope`.
+    pub fn requested_scope(&self, scope: &str) -> bool {
+        self.scopes.lock().unwrap().iter().any(|s| s == scope)
+    }
+}
+
+/// A [`TokenCredential`] that produces the emulator's fake AAD token.
+#[derive(Debug, Clone)]
+pub struct CosmosEmulatorCredential {
+    master_key: String,
+    call_count: Arc<AtomicUsize>,
+    scopes: Arc<Mutex<Vec<String>>>,
+}
+
+impl std::fmt::Debug for CredentialRecorder {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("CredentialRecorder")
+            .field("call_count", &self.call_count())
+            .finish_non_exhaustive()
+    }
+}
+
+impl CosmosEmulatorCredential {
+    /// Creates a credential using the emulator's well-known master key.
+    pub fn new() -> Self {
+        Self::with_master_key(EMULATOR_MASTER_KEY)
+    }
+
+    /// Creates a credential signing tokens with a specific master key string.
+    ///
+    /// The `master_key` must be the same key the emulator was started with; the
+    /// emulator validates the JWT signature segment against it byte-for-byte.
+    pub fn with_master_key(master_key: impl Into<String>) -> Self {
+        Self {
+            master_key: master_key.into(),
+            call_count: Arc::new(AtomicUsize::new(0)),
+            scopes: Arc::new(Mutex::new(Vec::new())),
+        }
+    }
+
+    /// Returns a [`CredentialRecorder`] that observes how this credential is used.
+    pub fn recorder(&self) -> CredentialRecorder {
+        CredentialRecorder {
+            call_count: self.call_count.clone(),
+            scopes: self.scopes.clone(),
+        }
+    }
+
+    /// Builds the emulator fake JWT: `b64url(header).b64url(claims).b64url(key)`.
+    fn build_token(&self, expires_on: OffsetDateTime) -> String {
+        // Header advertises the emulator's well-known signing key identifier.
+        let header = r#"{"typ":"JWT","alg":"RS256","x5t":"CosmosEmulatorPrimaryMaster","kid":"CosmosEmulatorPrimaryMaster"}"#;
+
+        let now = OffsetDateTime::now_utc().unix_timestamp();
+        // Allow for some clock skew between the test host and the emulator.
+        let not_before = now - 300;
+        let issued_at = now - 300;
+        let expiry = expires_on.unix_timestamp();
+
+        let claims = format!(
+            concat!(
+                "{{",
+                r#""aud":"https://localhost.localhost","#,
+                r#""iss":"https://sts.fake-issuer.net/7b1999a1-dfd7-440e-8204-00170979b984","#,
+                r#""iat":{iat},"nbf":{nbf},"exp":{exp},"#,
+                r#""aio":"","appid":"localhost","appidacr":"1","#,
+                r#""idp":"https://localhost.localhost","oid":"96313034-4739-43cb-93cd-74193adbe5b6","#,
+                r#""rh":"","sub":"localhost","#,
+                r#""tid":"EmulatorFederation","#,
+                r#""unique_name":"localhost","uti":"","#,
+                r#""ver":"1.0","#,
+                r#""scp":"user_impersonation","#,
+                r#""groups":["7ce1d3a1-6df6-4e88-bb0d-1c5e1c6c2b21","e99e9b3e-9c4f-4e5e-9e0a-2b9a4b2c3d4e","b2b3b4b5-c6c7-48d9-9e0a-2b9a4b2c3d4e","c3c4c5c6-d7d8-49e0-9f0b-3c0b5c3d4e5f","d4d5d6d7-e8e9-4af1-a01c-4d1c6d4e5f60"]"#,
+                "}}",
+            ),
+            iat = issued_at,
+            nbf = not_before,
+            exp = expiry,
+        );
+
+        let header_segment = URL_SAFE_NO_PAD.encode(header.as_bytes());
+        let claims_segment = URL_SAFE_NO_PAD.encode(claims.as_bytes());
+        // The emulator validates this segment against its master key string bytes.
+        let signature_segment = URL_SAFE_NO_PAD.encode(self.master_key.as_bytes());
+
+        format!("{header_segment}.{claims_segment}.{signature_segment}")
+    }
+}
+
+impl Default for CosmosEmulatorCredential {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
+impl TokenCredential for CosmosEmulatorCredential {
+    async fn get_token(
+        &self,
+        scopes: &[&str],
+        _options: Option<TokenRequestOptions<'_>>,
+    ) -> azure_core::Result<AccessToken> {
+        self.call_count.fetch_add(1, Ordering::SeqCst);
+        {
+            let mut recorded = self.scopes.lock().unwrap();
+            recorded.extend(scopes.iter().map(|s| s.to_string()));
+        }
+
+        let expires_on = OffsetDateTime::now_utc() + azure_core::time::Duration::hours(2);
+        let token = self.build_token(expires_on);
+        Ok(AccessToken::new(token, expires_on))
+    }
+}
diff --git a/sdk/cosmos/azure_data_cosmos/tests/framework/mod.rs b/sdk/cosmos/azure_data_cosmos/tests/framework/mod.rs
@@ -14,10 +14,12 @@
 //!
 //! The framework allows tests to easily run against real Cosmos DB instances, the local emulator, or a mock server using test-proxy.
 
+pub mod emulator_credential;
 pub mod mock_account;
 pub mod test_client;
 pub mod test_data;
 
+pub use emulator_credential::{CosmosEmulatorCredential, CredentialRecorder};
 pub use test_client::{
     assert_local_retry_attempted_on_region, assert_region_contacted_with_retry,
     get_effective_hub_endpoint, get_global_endpoint, resolve_connection_string, TestClient,
diff --git a/sdk/cosmos/azure_data_cosmos/tests/framework/test_client.rs b/sdk/cosmos/azure_data_cosmos/tests/framework/test_client.rs
diff --git a/sdk/cosmos/eng/scripts/Invoke-CosmosTestSetup.ps1 b/sdk/cosmos/eng/scripts/Invoke-CosmosTestSetup.ps1
