Replace Throughput Bugfix (#41564)

* replace throughput bugfix

* minor changes to tests

* added async and negative testing for replace throughput

* moved increment percent back into max_throughput if statement

* added pytest.raises to negative test

---------

Co-authored-by: Andrew Mathew <[email protected]>

Author: andrewmathew1

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -797,8 +797,8 @@ def _replace_throughput(
                 new_throughput_properties['content']['offerAutopilotSettings']['maxThroughput'] = max_throughput
                 if increment_percent:
                     new_throughput_properties['content']['offerAutopilotSettings']['autoUpgradePolicy']['throughputPolicy']['incrementPercent'] = increment_percent  # pylint: disable=line-too-long
-                if throughput.offer_throughput:
-                    new_throughput_properties["content"]["offerThroughput"] = throughput.offer_throughput
+            if throughput.offer_throughput:
+                new_throughput_properties["content"]["offerThroughput"] = throughput.offer_throughput
         except AttributeError as e:
             raise TypeError("offer_throughput must be int or an instance of ThroughputProperties") from e
 

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

@@ -24,7 +24,7 @@
 import azure.cosmos.documents as documents
 import azure.cosmos.exceptions as exceptions
 import test_config
-from azure.cosmos import _retry_utility
+from azure.cosmos import _retry_utility, ThroughputProperties
 from azure.cosmos.http_constants import HttpHeaders, StatusCodes
 from azure.cosmos.partition_key import PartitionKey
 
@@ -954,7 +954,34 @@ def test_index_progress_headers(self):
 
         created_db.delete_container(none_coll)
 
+    def test_replace_throughput_offer_with_int(self):
+        created_db = self.databaseForTest
+        collection = created_db.get_container_client(self.configs.TEST_MULTI_PARTITION_CONTAINER_ID)
+
+        new_throughput = ThroughputProperties(offer_throughput=2500)
+        collection.replace_throughput(new_throughput.offer_throughput)
+
+        retrieve_throughput = collection.get_throughput()
+        assert getattr(retrieve_throughput, "offer_throughput") == getattr(new_throughput, "offer_throughput")
+
+    def test_replace_throughput_offer_with_object(self):
+        created_db = self.databaseForTest
+        collection = created_db.get_container_client(self.configs.TEST_MULTI_PARTITION_CONTAINER_ID)
+
+        new_throughput = ThroughputProperties(offer_throughput=2500)
+        collection.replace_throughput(new_throughput)
+
+        retrieve_throughput = collection.get_throughput()
+        assert getattr(retrieve_throughput, "offer_throughput") == getattr(new_throughput, "offer_throughput")
+
+    def test_negative_replace_throughput_with_all_configs_set(self):
+        created_db = self.databaseForTest
+        collection = created_db.get_container_client(self.configs.TEST_MULTI_PARTITION_CONTAINER_ID)
+
+        new_throughput = ThroughputProperties(offer_throughput=2500, auto_scale_max_throughput=4000, auto_scale_increment_percent=5)
 
+        with pytest.raises(KeyError):
+            collection.replace_throughput(new_throughput)
 
     def _MockExecuteFunction(self, function, *args, **kwargs):
         if HttpHeaders.PartitionKey in args[4].headers:

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

@@ -25,6 +25,7 @@
 from azure.cosmos.aio import CosmosClient, _retry_utility_async, DatabaseProxy
 from azure.cosmos.http_constants import HttpHeaders, StatusCodes
 from azure.cosmos.partition_key import PartitionKey
+from azure.cosmos import ThroughputProperties
 
 
 class TimeoutTransport(AsyncioRequestsTransport):
@@ -956,7 +957,6 @@ def __validate_offer_response_body(self, offer, expected_coll_link, expected_off
             assert expected_offer_type == offer.properties.get('offerType')
 
     async def test_offer_read_and_query_async(self):
-
         # Create database.
         db = self.database_for_test
 
@@ -968,7 +968,6 @@ async def test_offer_read_and_query_async(self):
         self.__validate_offer_response_body(expected_offer, collection_properties.get('_self'), None)
 
     async def test_offer_replace_async(self):
-
         collection = self.database_for_test.get_container_client(self.configs.TEST_MULTI_PARTITION_CONTAINER_ID)
         # Read Offer
         expected_offer = await collection.get_throughput()
@@ -984,7 +983,6 @@ async def test_offer_replace_async(self):
         assert expected_offer.offer_throughput + 100 == replaced_offer.offer_throughput
 
     async def test_index_progress_headers_async(self):
-
         created_db = self.database_for_test
         consistent_coll = created_db.get_container_client(self.configs.TEST_MULTI_PARTITION_CONTAINER_ID)
         created_container = created_db.get_container_client(container=consistent_coll)
@@ -1007,6 +1005,35 @@ async def test_index_progress_headers_async(self):
 
         await created_db.delete_container(none_coll)
 
+    async def test_replace_throughput_offer_with_int(self):
+        created_db = self.database_for_test
+        collection = created_db.get_container_client(self.configs.TEST_MULTI_PARTITION_CONTAINER_ID)
+
+        new_throughput = ThroughputProperties(offer_throughput=2500)
+        await collection.replace_throughput(new_throughput.offer_throughput)
+
+        retrieve_throughput = await collection.get_throughput()
+        assert getattr(retrieve_throughput, "offer_throughput") == getattr(new_throughput, "offer_throughput")
+
+    async def test_replace_throughput_offer_with_object(self):
+        created_db = self.database_for_test
+        collection = created_db.get_container_client(self.configs.TEST_MULTI_PARTITION_CONTAINER_ID)
+
+        new_throughput = ThroughputProperties(offer_throughput=2500)
+        await collection.replace_throughput(new_throughput)
+
+        retrieve_throughput = await collection.get_throughput()
+        assert getattr(retrieve_throughput, "offer_throughput") == getattr(new_throughput, "offer_throughput")
+
+    async def test_negative_replace_throughput_with_all_configs_set(self):
+        created_db = self.database_for_test
+        collection = created_db.get_container_client(self.configs.TEST_MULTI_PARTITION_CONTAINER_ID)
+
+        new_throughput = ThroughputProperties(offer_throughput=2500, auto_scale_max_throughput=4000, auto_scale_increment_percent=5)
+
+        with pytest.raises(KeyError):
+            await collection.replace_throughput(new_throughput)
+
     async def _mock_execute_function(self, function, *args, **kwargs):
         if HttpHeaders.PartitionKey in args[4].headers:
             self.last_headers.append(args[4].headers[HttpHeaders.PartitionKey])