From 58f7ab0d2578d19e6061d37dcddc955c7edba9e9 Mon Sep 17 00:00:00 2001
From: rajinder <rajendersaini@gmail.com>
Date: Wed, 19 May 2021 23:10:08 -0700
Subject: [PATCH] Refector Apigateway update_usage_plan and add all replace
 operation. (#3942)

* Refector Apigateway update_usage_plan and add all replace operation.

* code reformatting

Co-authored-by: rajinder saini <rajinder.saini@c02vt5k2htd6.corp.climate.com>
---
 moto/apigateway/models.py                | 38 ++++++++++++++----------
 tests/test_apigateway/test_apigateway.py |  7 ++++-
 2 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/moto/apigateway/models.py b/moto/apigateway/models.py
index eea3cbfeb..45fc1e8c9 100644
--- a/moto/apigateway/models.py
+++ b/moto/apigateway/models.py
@@ -545,6 +545,7 @@ class UsagePlan(BaseModel, dict):
         apiStages=None,
         throttle=None,
         quota=None,
+        productCode=None,
         tags=None,
     ):
         super(UsagePlan, self).__init__()
@@ -554,8 +555,29 @@ class UsagePlan(BaseModel, dict):
         self["apiStages"] = apiStages if apiStages else []
         self["throttle"] = throttle
         self["quota"] = quota
+        self["productCode"] = productCode
         self["tags"] = tags
 
+    def apply_patch_operations(self, patch_operations):
+        for op in patch_operations:
+            path = op["path"]
+            value = op["value"]
+            if op["op"] == "replace":
+                if "/name" in path:
+                    self["name"] = value
+                if "/productCode" in path:
+                    self["productCode"] = value
+                if "/description" in path:
+                    self["description"] = value
+                if "/quota/limit" in path:
+                    self["quota"]["limit"] = value
+                if "/quota/period" in path:
+                    self["quota"]["period"] = value
+                if "/throttle/rateLimit" in path:
+                    self["throttle"]["rateLimit"] = value
+                if "/throttle/burstLimit" in path:
+                    self["throttle"]["burstLimit"] = value
+
 
 class UsagePlanKey(BaseModel, dict):
     def __init__(self, id, type, name, value):
@@ -1212,26 +1234,12 @@ class APIGatewayBackend(BaseBackend):
     def get_usage_plan(self, usage_plan_id):
         if usage_plan_id not in self.usage_plans:
             raise UsagePlanNotFoundException()
-
         return self.usage_plans[usage_plan_id]
 
-    def __apply_usage_plan_patch_operations(self, plan, patch_operations):
-        for op in patch_operations:
-            if op["op"] == "replace":
-                if "/quota/limit" in op["path"]:
-                    plan["quota"]["limit"] = op["value"]
-                if "/quota/period" in op["path"]:
-                    plan["quota"]["period"] = op["value"]
-                if "/throttle/rateLimit" in op["path"]:
-                    plan["throttle"]["rateLimit"] = op["value"]
-                if "/throttle/burstLimit" in op["path"]:
-                    plan["throttle"]["burstLimit"] = op["value"]
-
     def update_usage_plan(self, usage_plan_id, patch_operations):
         if usage_plan_id not in self.usage_plans:
             raise UsagePlanNotFoundException()
-        plan = self.usage_plans[usage_plan_id]
-        self.__apply_usage_plan_patch_operations(plan, patch_operations)
+        self.usage_plans[usage_plan_id].apply_patch_operations(patch_operations)
         return self.usage_plans[usage_plan_id]
 
     def delete_usage_plan(self, usage_plan_id):
diff --git a/tests/test_apigateway/test_apigateway.py b/tests/test_apigateway/test_apigateway.py
index 090e30f9b..f7e7f4180 100644
--- a/tests/test_apigateway/test_apigateway.py
+++ b/tests/test_apigateway/test_apigateway.py
@@ -2081,12 +2081,17 @@ def test_update_usage_plan():
             {"op": "replace", "path": "/quota/period", "value": "MONTH"},
             {"op": "replace", "path": "/throttle/rateLimit", "value": "500"},
             {"op": "replace", "path": "/throttle/burstLimit", "value": "1500"},
+            {"op": "replace", "path": "/name", "value": "new-name"},
+            {"op": "replace", "path": "/description", "value": "new-description"},
+            {"op": "replace", "path": "/productCode", "value": "new-productionCode"},
         ],
     )
     response["quota"]["limit"].should.equal("1000")
     response["quota"]["period"].should.equal("MONTH")
     response["quota"]["limit"].should.equal("1000")
-    response["quota"]["limit"].should.equal("1000")
+    response["name"].should.equal("new-name")
+    response["description"].should.equal("new-description")
+    response["productCode"].should.equal("new-productionCode")
 
 
 @mock_apigateway