AAD test coverage (#46568)

* AAD test coverage initial commit

* commenting all AAD coverage except one as a smoke test

* adding pipeline changes

* adding documentation

* update the naming convention such that our automation will collect it automagically

* fixing CI pipeline binding

* expanding coverage of AAD tests

* expanding coverage of AAD tests

* cleaning up resources

* cleaning up resources

* cleaning up resources

* fixing pipeline errors

* debugging pipleine

* debugging pipleine

* running a subset

* enabling AAD at a bigger scale

* AAD fixing tests

* AAD fixing tests

* fixing AAD tests

* AAD fixing tests

* AAD fixing tests

* fixing AAD tests

* addressing PR comments

* fixing bugs

* fixing Copilot comments

* fixing Copilot comments

* refactoring tests

---------

Co-authored-by: Scott Beddall <scbedd@microsoft.com>

Author: dibahlfi

sdk/cosmos/azure-cosmos/pytest.ini

@@ -3,6 +3,13 @@ markers =
     cosmosEmulator: marks tests as depending in Cosmos DB Emulator.
     cosmosLong: marks tests to be run on a Cosmos DB live account.
     cosmosQuery: marks tests running queries on Cosmos DB live account.
+    cosmosAADLong: marks AAD tests for the standard live-account lane.
+    cosmosAADSplit: marks AAD tests for partition split scenarios.
+    cosmosAADMultiRegion: marks AAD tests for multi-region scenarios.
+    cosmosAADCircuitBreaker: marks AAD tests for circuit-breaker scenarios.
+    cosmosAADCircuitBreakerMultiRegion: marks AAD tests for single-master multi-region-read circuit-breaker scenarios.
+    cosmosAADQuery: marks AAD tests for query-focused scenarios.
+    cosmosAADPerPartitionAutomaticFailover: marks AAD tests for per-partition automatic failover scenarios.
     cosmosSplit: marks test where there are partition splits on CosmosDB live account.
     cosmosMultiRegion: marks tests running on a Cosmos DB live account with multi-region and multi-write enabled.
     cosmosCircuitBreaker: marks tests running on Cosmos DB live account with per partition circuit breaker enabled and multi-write enabled.

sdk/cosmos/azure-cosmos/tests/test_aad.py

@@ -13,9 +13,11 @@
 
 import azure.cosmos.cosmos_client as cosmos_client
 import test_config
-from azure.cosmos import DatabaseProxy, ContainerProxy, exceptions
+from azure.cosmos import DatabaseProxy, ContainerProxy
 from azure.core.exceptions import HttpResponseError
 
+
+
 def _remove_padding(encoded_string):
     while encoded_string.endswith("="):
         encoded_string = encoded_string[0:len(encoded_string) - 1]
@@ -35,7 +37,7 @@ def get_test_item(num):
 
 class CosmosEmulatorCredential(object):
     def get_token(self, *scopes, **kwargs):
-        # type: (*str, **Any) -> AccessToken
+        # type: (*str, **object) -> AccessToken
         """Request an access token for the emulator. Based on Azure Core's Access Token Credential.
 
         This method is called automatically by Azure SDK clients.
@@ -93,14 +95,21 @@ class TestAAD(unittest.TestCase):
     configs = test_config.TestConfig
     host = configs.host
     masterKey = configs.masterKey
-    credential = CosmosEmulatorCredential() if configs.is_emulator else configs.credential
+    # Emulator-only credential used by this class.
+    credential = CosmosEmulatorCredential()
+    _skip_on_non_emulator = pytest.mark.skipif(
+        not configs.is_emulator,
+        reason="Emulator credential tests are emulator-specific (localhost audience)."
+    )
 
     @classmethod
     def setUpClass(cls):
+        # Emulator-only path: always use the emulator credential.
         cls.client = cosmos_client.CosmosClient(cls.host, cls.credential)
         cls.database = cls.client.get_database_client(cls.configs.TEST_DATABASE_ID)
         cls.container = cls.database.get_container_client(cls.configs.TEST_SINGLE_PARTITION_CONTAINER_ID)
 
+    @_skip_on_non_emulator
     def test_aad_credentials(self):
         print("Container info: " + str(self.container.read()))
         self.container.create_item(get_test_item(0))
@@ -110,14 +119,6 @@ def test_aad_credentials(self):
         print("Query result: " + str(query_results[0]))
         self.container.delete_item(item='Item_0', partition_key='pk')
 
-        # Attempting to do management operations will return a 403 Forbidden exception
-        try:
-            self.client.delete_database(self.configs.TEST_DATABASE_ID)
-        except exceptions.CosmosHttpResponseError as e:
-            assert e.status_code == 403
-            print("403 error assertion success")
-
-
     def _run_with_scope_capture(self, credential_cls, action, *args, **kwargs):
         scopes_captured = []
         original_get_token = credential_cls.get_token
@@ -133,6 +134,7 @@ def capturing_get_token(self, *scopes, **kwargs):
             credential_cls.get_token = original_get_token
         return scopes_captured, result
 
+    @_skip_on_non_emulator
     def test_override_scope_no_fallback(self):
         """When override scope is provided, only that scope is used and no fallback occurs."""
         override_scope = "https://my.custom.scope/.default"
@@ -156,6 +158,7 @@ def action(scopes_captured):
             except Exception:
                 pass
 
+    @_skip_on_non_emulator
     def test_override_scope_auth_error_no_fallback(self):
         """When override scope is provided and auth fails, no fallback to other scopes occurs."""
         override_scope = "https://my.custom.scope/.default"
@@ -180,6 +183,7 @@ def action(scopes_captured):
         finally:
             del os.environ["AZURE_COSMOS_AAD_SCOPE_OVERRIDE"]
 
+    @_skip_on_non_emulator
     def test_account_scope_only(self):
         """When account scope is provided, only that scope is used."""
         account_scope = "https://localhost/.default"
@@ -203,6 +207,7 @@ def action(scopes_captured):
             except Exception:
                 pass
 
+    @_skip_on_non_emulator
     def test_account_scope_fallback_on_error(self):
         """When account scope is provided and auth fails, fallback to default scope occurs."""
         account_scope = "https://localhost/.default"

sdk/cosmos/azure-cosmos/tests/test_aad_async.py

@@ -12,10 +12,11 @@
 from azure.core.credentials import AccessToken
 
 import test_config
-from azure.cosmos import exceptions
 from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy
 from azure.core.exceptions import HttpResponseError
 
+
+
 def _remove_padding(encoded_string):
     while encoded_string.endswith("="):
         encoded_string = encoded_string[0:len(encoded_string) - 1]
@@ -35,7 +36,7 @@ def get_test_item(num):
 
 class CosmosEmulatorCredential(object):
     async def get_token(self, *scopes, **kwargs):
-        # type: (*str, **Any) -> AccessToken
+        # type: (*str, **object) -> AccessToken
         """Request an access token for the emulator. Based on Azure Core's Access Token Credential.
 
         This method is called automatically by Azure SDK clients.
@@ -93,16 +94,11 @@ class TestAADAsync(unittest.IsolatedAsyncioTestCase):
     configs = test_config.TestConfig
     host = configs.host
     masterKey = configs.masterKey
-    credential = CosmosEmulatorCredential() if configs.is_emulator else configs.credential_async
-
-    @classmethod
-    def setUpClass(cls):
-        if (cls.credential == '[YOUR_KEY_HERE]' or
-                cls.host == '[YOUR_ENDPOINT_HERE]'):
-            raise Exception(
-                "You must specify your Azure Cosmos account values for "
-                "'masterKey' and 'host' at the top of this class to run the "
-                "tests.")
+    credential = CosmosEmulatorCredential()
+    _skip_scope_tests_on_non_emulator = pytest.mark.skipif(
+        not configs.is_emulator,
+        reason="Scope capture tests are emulator-specific (localhost audience)."
+    )
 
     async def asyncSetUp(self):
         self.client = CosmosClient(self.host, self.credential)
@@ -112,9 +108,8 @@ async def asyncSetUp(self):
     async def asyncTearDown(self):
         await self.client.close()
 
+    @_skip_scope_tests_on_non_emulator
     async def test_aad_credentials_async(self):
-        # Do any R/W data operations with your authorized AAD client
-
         print("Container info: " + str(await self.container.read()))
         await self.container.create_item(get_test_item(0))
         print("Point read result: " + str(await self.container.read_item(item='Item_0', partition_key='pk')))
@@ -123,12 +118,6 @@ async def test_aad_credentials_async(self):
         print("Query result: " + str(query_results[0]))
         await self.container.delete_item(item='Item_0', partition_key='pk')
 
-        # Attempting to do management operations will return a 403 Forbidden exception
-        try:
-            await self.client.delete_database(self.configs.TEST_DATABASE_ID)
-        except exceptions.CosmosHttpResponseError as e:
-            assert e.status_code == 403
-            print("403 error assertion success")
 
     async def _run_with_scope_capture_async(self, credential_cls, action):
         scopes_captured = []
@@ -146,6 +135,7 @@ async def capturing_get_token(self, *scopes, **kwargs):
         finally:
             credential_cls.get_token = orig_get_token
 
+    @_skip_scope_tests_on_non_emulator
     async def test_override_scope_no_fallback_async(self):
         """When override scope is provided, only that scope is used and no fallback occurs."""
         override_scope = "https://my.custom.scope/.default"
@@ -172,6 +162,7 @@ async def action(scopes_captured):
             except Exception:
                 pass
 
+    @_skip_scope_tests_on_non_emulator
     async def test_override_scope_no_fallback_on_error_async(self):
         """When override scope is provided and auth fails, no fallback occurs."""
         override_scope = "https://my.custom.scope/.default"
@@ -205,6 +196,7 @@ async def action(scopes_captured):
             except Exception:
                 pass
 
+    @_skip_scope_tests_on_non_emulator
     async def test_account_scope_only_async(self):
         """When account scope is provided, only that scope is used."""
         account_scope = "https://localhost/.default"
@@ -230,6 +222,7 @@ async def action(scopes_captured):
             except Exception:
                 pass
 
+    @_skip_scope_tests_on_non_emulator
     async def test_account_scope_fallback_on_error_async(self):
         """When account scope is provided and auth fails, fallback to default scope occurs."""
         account_scope = "https://localhost/.default"

sdk/cosmos/azure-cosmos/tests/test_aggregate.py

@@ -15,21 +15,24 @@
 
 
 class _config:
+    is_aad_mode = test_config.TestConfig.data_auth_mode == "aad"
     host = test_config.TestConfig.host
     master_key = test_config.TestConfig.masterKey
     connection_policy = test_config.TestConfig.connectionPolicy
     PARTITION_KEY = 'key'
     UNIQUE_PARTITION_KEY = 'uniquePartitionKey'
     FIELD = 'field'
-    DOCUMENTS_COUNT = 400
-    DOCS_WITH_SAME_PARTITION_KEY = 200
+    # Keep key-auth query coverage unchanged; trim only AAD runs to stay under CI timeout.
+    DOCUMENTS_COUNT = 120 if is_aad_mode else 400
+    DOCS_WITH_SAME_PARTITION_KEY = 60 if is_aad_mode else 200
     docs_with_numeric_id = 0
     sum = 0
 
 
 @pytest.mark.cosmosQuery
 class TestAggregateQuery(unittest.TestCase):
     client: cosmos_client.CosmosClient = None
+    key_client: cosmos_client.CosmosClient = None
 
     @classmethod
     def setUpClass(cls):
@@ -40,7 +43,7 @@ def setUpClass(cls):
     @classmethod
     def tearDownClass(cls) -> None:
         try:
-            cls.created_db.delete_container(cls.created_collection.id)
+            cls.key_db.delete_container(cls.created_collection.id)
         except CosmosHttpResponseError:
             pass
 
@@ -52,9 +55,10 @@ def _setup(cls):
                 "'masterKey' and 'host' at the top of this class to run the "
                 "tests.")
 
-        cls.client = cosmos_client.CosmosClient(_config.host, _config.master_key)
-        cls.created_db = cls.client.get_database_client(test_config.TestConfig.TEST_DATABASE_ID)
-        cls.created_collection = cls._create_collection(cls.created_db)
+        cls.key_client, cls.key_db, cls.client, cls.created_db = (
+            test_config.TestConfig.create_test_clients(test_config.TestConfig.TEST_DATABASE_ID))
+        created_collection_ref = cls._create_collection(cls.key_db)
+        cls.created_collection = cls.created_db.get_container_client(created_collection_ref.id)
 
         # test documents
         document_definitions = []
@@ -138,6 +142,62 @@ def test_run_all(self):
                 print(test_name + ': ' + query + " FAILED")
                 raise e
 
+    # AAD-only smoke subset.
+    #
+    # Why this exists: the CI AAD lane runs on Linux and the shared
+    # ``azpysdk.main whl --isolate`` bootstrap on that pool already eats
+    # ~90 minutes of the 120-minute job ceiling. Running the full
+    # ``test_run_all`` matrix (24 aggregate variants) under AAD on top of
+    # that bootstrap pushes the lane over the ceiling. The full matrix
+    # still runs under the ``cosmosQuery`` lane (key auth) -- this method
+    # is *additional* AAD-only coverage focused on Contoso's exact bug
+    # shape: cross-partition aggregate query under bearer auth, including
+    # the ORDER BY pagination case where token refresh mid-stream is most
+    # likely to surface.
+    #
+    # Three queries: cross-partition COUNT (fan-out), cross-partition SUM
+    # with ORDER BY (fan-out + paginated reduce -> token-refresh window),
+    # single-partition AVG (pinned-PK path).
+    @pytest.mark.cosmosAADLong
+    @pytest.mark.skipif(
+        test_config.TestConfig.data_auth_mode != "aad",
+        reason="AAD-only smoke subset; full coverage runs under cosmosQuery (key auth).",
+    )
+    def test_aad_aggregate_subset(self):
+        same_partition_avg = (
+            _config.DOCS_WITH_SAME_PARTITION_KEY * (_config.DOCS_WITH_SAME_PARTITION_KEY + 1) / 2.0
+        ) / _config.DOCS_WITH_SAME_PARTITION_KEY
+        subset = [
+            (
+                "test_aad_xp_count",
+                "SELECT VALUE COUNT(r.{}) FROM r WHERE true".format(_config.PARTITION_KEY),
+                _config.DOCUMENTS_COUNT,
+            ),
+            (
+                "test_aad_xp_sum_orderby",
+                "SELECT VALUE SUM(r.{f}) FROM r WHERE IS_NUMBER(r.{pk}) ORDER BY r.{pk}".format(
+                    f=_config.PARTITION_KEY, pk=_config.PARTITION_KEY
+                ),
+                _config.sum,
+            ),
+            (
+                "test_aad_sp_avg",
+                "SELECT VALUE AVG(r.{f}) FROM r WHERE r.{pk} = '{val}'".format(
+                    f=_config.FIELD,
+                    pk=_config.PARTITION_KEY,
+                    val=_config.UNIQUE_PARTITION_KEY,
+                ),
+                same_partition_avg,
+            ),
+        ]
+        for test_name, query, expected in subset:
+            try:
+                self._run_one(query, expected)
+                print(test_name + ': ' + query + " PASSED", flush=True)
+            except Exception as e:
+                print(test_name + ': ' + query + " FAILED", flush=True)
+                raise e
+
     def _run_one(self, query, expected_result):
         self._execute_query_and_validate_results(self.created_collection, query, expected_result)
 

sdk/cosmos/azure-cosmos/tests/test_auto_scale.py

@@ -13,7 +13,7 @@
 
 @pytest.mark.cosmosLong
 class TestAutoScale(unittest.TestCase):
-    client: CosmosClient = None
+    key_client: CosmosClient = None
     host = test_config.TestConfig.host
     masterKey = test_config.TestConfig.masterKey
     connectionPolicy = test_config.TestConfig.connectionPolicy
@@ -27,8 +27,8 @@ def setUpClass(cls):
                 "'masterKey' and 'host' at the top of this class to run the "
                 "tests.")
 
-        cls.client = cosmos_client.CosmosClient(cls.host, cls.masterKey)
-        cls.created_database = cls.client.get_database_client(test_config.TestConfig.TEST_DATABASE_ID)
+        cls.key_client = cosmos_client.CosmosClient(cls.host, cls.masterKey)
+        cls.created_database = cls.key_client.get_database_client(test_config.TestConfig.TEST_DATABASE_ID)
 
     def test_autoscale_create_container(self):
         container_id = None
@@ -75,7 +75,7 @@ def test_autoscale_create_database(self):
         database_id = "db_auto_scale_" + str(uuid.uuid4())
         try:
             # Testing auto_scale_settings for the create_database method
-            created_database = self.client.create_database(database_id, offer_throughput=ThroughputProperties(
+            created_database = self.key_client.create_database(database_id, offer_throughput=ThroughputProperties(
                 auto_scale_max_throughput=5000,
                 auto_scale_increment_percent=2))
             created_db_properties = created_database.get_throughput()
@@ -84,11 +84,11 @@ def test_autoscale_create_database(self):
             # Testing the input value of the increment_percentage
             assert created_db_properties.auto_scale_increment_percent == 2
 
-            self.client.delete_database(created_database.id)
+            self.key_client.delete_database(created_database.id)
 
             # Testing auto_scale_settings for the create_database_if_not_exists method
             database_id = "db_auto_scale_2_" + str(uuid.uuid4())
-            created_database = self.client.create_database_if_not_exists(database_id,
+            created_database = self.key_client.create_database_if_not_exists(database_id,
                                                                          offer_throughput=ThroughputProperties(
                                                                              auto_scale_max_throughput=9000,
                                                                              auto_scale_increment_percent=11))
@@ -98,13 +98,13 @@ def test_autoscale_create_database(self):
             # Testing the input value of the increment_percentage
             assert created_db_properties.auto_scale_increment_percent == 11
         finally:
-            self.client.delete_database(database_id)
+            self.key_client.delete_database(database_id)
 
     def test_autoscale_replace_throughput(self):
         database_id = "replace_db" + str(uuid.uuid4())
         container_id = None
         try:
-            created_database = self.client.create_database(database_id, offer_throughput=ThroughputProperties(
+            created_database = self.key_client.create_database(database_id, offer_throughput=ThroughputProperties(
                 auto_scale_max_throughput=5000,
                 auto_scale_increment_percent=2))
             created_database.replace_throughput(
@@ -114,7 +114,7 @@ def test_autoscale_replace_throughput(self):
             assert created_db_properties.auto_scale_max_throughput == 7000
             # Testing the input value of the increment_percentage
             assert created_db_properties.auto_scale_increment_percent == 20
-            self.client.delete_database(database_id)
+            self.key_client.delete_database(database_id)
 
             container_id = "container_with_auto_scale_settings" + str(uuid.uuid4())
             created_container = self.created_database.create_container(