Update python unit tests

caraml-dev · May 23, 2024 · 4104721 · 4104721
1 parent 360729b
commit 4104721
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 7 deletions.
diff --git a/sdk/tests/conftest.py b/sdk/tests/conftest.py
@@ -277,7 +277,7 @@ def generic_router_version_status():
 @pytest.fixture
 def generic_resource_request():
     return turing.generated.models.ResourceRequest(
-        min_replica=1, max_replica=3, cpu_request="100m", memory_request="512Mi"
+        min_replica=1, max_replica=3, cpu_request="100m", memory_request="512Mi", cpu_limit=None,
     )
 
 

diff --git a/sdk/tests/router/config/router_ensembler_config_test.py b/sdk/tests/router/config/router_ensembler_config_test.py
@@ -600,7 +600,7 @@ def test_create_nop_router_ensembler_config_with_invalid_route(
                     image="test.io/just-a-test/turing-ensembler:0.0.0-build.0",
                     port=5120,
                     resource_request=turing.generated.models.ResourceRequest(
-                        cpu_request="100m", max_replica=3,
+                        cpu_request="100m", cpu_limit=None, max_replica=3,
                         memory_request="512Mi", min_replica=1,
                     ),
                     service_account="secret-name-for-google-service-account",
@@ -619,8 +619,8 @@ def test_create_nop_router_ensembler_config_with_invalid_route(
                     env=[turing.generated.models.EnvVar(name="env_name", value="env_val")],
                     project_id=77,
                     resource_request=turing.generated.models.ResourceRequest(
-                        cpu_request="100m", max_replica=3,
-                        memory_request="512Mi",min_replica=1,
+                        cpu_request="100m", cpu_limit=None, max_replica=3,
+                        memory_request="512Mi", min_replica=1,
                     ),
                     timeout="500ms"
                 ),

diff --git a/sdk/turing/generated/model/resource_request.py b/sdk/turing/generated/model/resource_request.py
@@ -90,7 +90,7 @@ def openapi_types():
             'min_replica': (int,),  # noqa: E501
             'max_replica': (int,),  # noqa: E501
             'cpu_request': (str,),  # noqa: E501
-            'cpu_limit': (str,),  # noqa: E501
+            'cpu_limit': (str, none_type,),  # noqa: E501
             'memory_request': (str,),  # noqa: E501
         }
 
@@ -156,7 +156,7 @@ def __init__(self, *args, **kwargs):  # noqa: E501
             min_replica (int): [optional]  # noqa: E501
             max_replica (int): [optional]  # noqa: E501
             cpu_request (str): [optional]  # noqa: E501
-            cpu_limit (str): [optional]  # noqa: E501
+            cpu_limit (str, none_type): [optional]  # noqa: E501
             memory_request (str): [optional]  # noqa: E501
         """
 

diff --git a/sdk/turing/router/config/resource_request.py b/sdk/turing/router/config/resource_request.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import ClassVar
+from typing import Optional
 
 import turing.generated.models
 from turing.generated.model_utils import OpenApiModel
@@ -11,11 +11,13 @@ class ResourceRequest:
     max_replica: int
     cpu_request: str
     memory_request: str
+    cpu_limit: Optional[str] = None
 
     _min_replica: int = field(init=False, repr=False)
     _max_replica: int = field(init=False, repr=False)
     _cpu_request: str = field(init=False, repr=False)
     _memory_request: str = field(init=False, repr=False)
+    _cpu_limit: Optional[str] = field(init=False, repr=False, default=None)
 
     @property
     def min_replica(self) -> int:
@@ -45,6 +47,16 @@ def cpu_request(self) -> str:
     def cpu_request(self, cpu_request: str):
         self._cpu_request = cpu_request
 
+    @property
+    def cpu_limit(self) -> str:
+        return self._cpu_limit
+
+    @cpu_limit.setter
+    def cpu_limit(self, cpu_limit):
+        if type(cpu_limit) is property:
+            cpu_limit = ResourceRequest._cpu_limit
+        self._cpu_limit = cpu_limit
+
     @property
     def memory_request(self) -> str:
         return self._memory_request
@@ -66,6 +78,7 @@ def to_open_api(self) -> OpenApiModel:
             min_replica=self.min_replica,
             max_replica=self.max_replica,
             cpu_request=self.cpu_request,
+            cpu_limit=self.cpu_limit,
             memory_request=self.memory_request,
         )