diff --git a/kr8s/_objects.py b/kr8s/_objects.py index d7ccbed..d4c026d 100644 --- a/kr8s/_objects.py +++ b/kr8s/_objects.py @@ -1228,6 +1228,43 @@ def gen( ) ) + async def tolerate( + self, + key: str, + *, + operator: str, + effect: str, + value: str | None = None, + toleration_seconds: int | None = None, + ): + """Add a toleration to the Pod. + + Args: + key (str): Key of the taint to tolerate. + operator (str): Toleration operator. + effect (str): Specifies the effect of the taint to tolerate. + value (str): Value of taint to tolerate. + toleration_seconds (str): Toleration seconds. + """ + new_toleration: dict = {"key": key, "operator": operator, "effect": effect} + if value is not None: + new_toleration["value"] = value + + if toleration_seconds is not None: + new_toleration["tolerationSeconds"] = toleration_seconds + + self.tolerations.append(new_toleration) + + await self.async_patch({"spec": {"tolerations": self.tolerations}}) + + @property + def tolerations(self) -> Box: + """Tolerations for Pod.""" + try: + return self.raw["spec"]["tolerations"] + except KeyError: + return Box([]) + class PodTemplate(APIObject): """A Kubernetes PodTemplate.""" diff --git a/kr8s/tests/test_objects.py b/kr8s/tests/test_objects.py index 25dafcd..7e9a814 100644 --- a/kr8s/tests/test_objects.py +++ b/kr8s/tests/test_objects.py @@ -350,6 +350,42 @@ def test_pod_get_sync(example_pod_spec): pod2.delete() +@pytest.fixture +async def example_pod(example_pod_spec): + pod = await Pod(example_pod_spec) + await pod.create() + yield pod + await pod.delete() + + +async def test_pod_tolerate(example_pod): + await example_pod.tolerate("key1", operator="Exists", effect="NoSchedule") + await example_pod.tolerate( + "key2", + operator="Equal", + effect="NoExecute", + value="value1", + toleration_seconds=600, + ) + pod2 = await Pod.get(example_pod.name, namespace=example_pod.namespace) + + pod_objects_to_check = [example_pod, pod2] + for pod in pod_objects_to_check: + key_to_toleration = {} + for toleration in pod.tolerations: + if toleration["key"] in ("key1", "key2"): + key_to_toleration[toleration["key"]] = toleration + + assert len(key_to_toleration) == 2 + assert key_to_toleration["key1"]["operator"] == "Exists" + assert key_to_toleration["key1"]["effect"] == "NoSchedule" + + assert key_to_toleration["key2"]["operator"] == "Equal" + assert key_to_toleration["key2"]["effect"] == "NoExecute" + assert key_to_toleration["key2"]["value"] == "value1" + assert key_to_toleration["key2"]["tolerationSeconds"] == 600 + + async def test_pod_from_name(example_pod_spec): pod = await Pod(example_pod_spec) await pod.create()