{cosmosdb-preview} Fix SQL container throughput update to preserve existing throughput buckets (Azure#9033)

* test: Enhance CosmosDB throughput bucketing tests to verify bucket retention and removal

* Changes to preserve throughput buckets on throughput change and record test

* Bump version in setup.py and fix whitespaces

* Fix formatting in error messages for container not found exceptions

* Rerecorded failing tests

Author: Achint-Agrawal

src/cosmosdb-preview/HISTORY.rst

@@ -1,6 +1,11 @@
 .. :changelog:
 Release History
 ===============
+
+1.6.1
++++++
+* Fix SQL container throughput update to preserve existing throughput buckets when not explicitly specified.
+
 1.6.0
 * Add support for Fleet/Fleetspace/FleetspaceAccount/FleetAnalytics CRUD actions.
 

src/cosmosdb-preview/azext_cosmosdb_preview/custom.py

@@ -15,14 +15,16 @@ def test_cosmosdb_sql_throughput_bucketing(self, resource_group):
         col = self.create_random_name(prefix='cli', length=15)
         db_name = self.create_random_name(prefix='cli', length=15)
         throughput_buckets = '"[{\\"id\\": 1, \\"maxThroughputPercentage\\": 10 }, {\\"id\\": 2, \\"maxThroughputPercentage\\": 20 }]"'
+        empty_throughput_buckets = '"[]"'
 
         # This test needs to be run on the following account since it has the feature flag enabled
         self.kwargs.update({
             'rg': 'throughput-bucketing-rg',
             'acc': 'throughput-bucketing-rp-test',
             'db_name': db_name,
             'col': col,
-            'throughput_buckets': throughput_buckets
+            'throughput_buckets': throughput_buckets,
+            'empty_throughput_buckets': empty_throughput_buckets
         })
 
         # Create database
@@ -44,6 +46,30 @@ def test_cosmosdb_sql_throughput_bucketing(self, resource_group):
         assert actual_throughput_buckets[1]["id"] == 2
         assert actual_throughput_buckets[1]["maxThroughputPercentage"] == 20
 
+        # Test: Change throughput without specifying throughput buckets - buckets should be retained
+        self.cmd('az cosmosdb sql container throughput update -g {rg} -a {acc} -d {db_name} -n {col} --throughput 4000')
+
+        throughput_resource_json = self.cmd('az cosmosdb sql container throughput show -g {rg} -a {acc} -d {db_name} -n {col}').get_output_in_json()
+        retained_throughput_buckets = throughput_resource_json["resource"]["throughputBuckets"]
+
+        # Verify throughput buckets are retained after throughput change
+        assert len(retained_throughput_buckets) == 2
+        assert retained_throughput_buckets[0]["id"] == 1
+        assert retained_throughput_buckets[0]["maxThroughputPercentage"] == 10
+        assert retained_throughput_buckets[1]["id"] == 2
+        assert retained_throughput_buckets[1]["maxThroughputPercentage"] == 20
+
+        # Verify throughput was actually updated
+        assert throughput_resource_json["resource"]["throughput"] == 4000
+
+        # Test: Set throughput buckets to empty array - buckets should be removed
+        self.cmd('az cosmosdb sql container throughput update -g {rg} -a {acc} -d {db_name} -n {col} --throughput 4000 --throughput-buckets {empty_throughput_buckets}')
+
+        throughput_resource_json = self.cmd('az cosmosdb sql container throughput show -g {rg} -a {acc} -d {db_name} -n {col}').get_output_in_json()
+        
+        # Verify throughput buckets are removed when set to empty array
+        assert len(throughput_resource_json["resource"]["throughputBuckets"]) == 0
+
         # delete container
         self.cmd('az cosmosdb sql container delete -g {rg} -a {acc} -d {db_name} -n {col} -y')
 

src/cosmosdb-preview/azext_cosmosdb_preview/tests/latest/recordings/test_cosmosdb_sql_container_adaptiveru.yaml

@@ -16,7 +16,7 @@
 
 # TODO: Confirm this is the right version number you want and it matches your
 # HISTORY.rst entry.
-VERSION = '1.6.0'
+VERSION = '1.6.1'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers