Adds tolerate to Pod (#538)

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."""

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()