From 537de43e4732c2a473b015c8c9f658513d2b5ec0 Mon Sep 17 00:00:00 2001
From: "Thomas J. Fan" <thomasjpfan@gmail.com>
Date: Tue, 17 Dec 2024 06:15:43 -0500
Subject: [PATCH] Add taint function and taints property to Node (#530)

Co-authored-by: Jacob Tomlinson <jacobtomlinson@users.noreply.github.com>
---
 kr8s/_objects.py           | 29 ++++++++++++++++++++++++++++-
 kr8s/tests/test_objects.py | 24 ++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/kr8s/_objects.py b/kr8s/_objects.py
index e2b1d8c..d7ccbed 100644
--- a/kr8s/_objects.py
+++ b/kr8s/_objects.py
@@ -383,7 +383,9 @@ async def patch(self, patch, *, subresource=None, type=None) -> None:
         """Patch this object in Kubernetes."""
         await self.async_patch(patch, subresource=subresource, type=type)
 
-    async def async_patch(self, patch: dict, *, subresource=None, type=None) -> None:
+    async def async_patch(
+        self, patch: dict | list, *, subresource=None, type=None
+    ) -> None:
         """Patch this object in Kubernetes."""
         url = f"{self.endpoint}/{self.name}"
         if type == "json":
@@ -828,6 +830,31 @@ async def uncordon(self) -> None:
         """
         await self.async_patch({"spec": {"unschedulable": False}})
 
+    async def taint(self, key: str, value: str, *, effect: str) -> None:
+        """Taint a node."""
+        await self.async_refresh()
+        if effect.endswith("-"):
+            # Remove taint with key
+            effect = effect[:-1]
+            if all(taint["key"] != key for taint in self.taints):
+                raise NotFoundError(f"Unable to find taint with key: {key}")
+
+            taints = [taint for taint in self.taints if taint["key"] != key]
+        else:
+            taints = list(self.taints) + [
+                {"key": key, "value": value, "effect": effect}
+            ]
+
+        await self.async_patch({"spec": {"taints": taints}})
+
+    @property
+    def taints(self) -> Box:
+        """Labels of the Kubernetes resource."""
+        try:
+            return self.raw["spec"]["taints"]
+        except KeyError:
+            return Box({})
+
 
 class PersistentVolumeClaim(APIObject):
     """A Kubernetes PersistentVolumeClaim."""
diff --git a/kr8s/tests/test_objects.py b/kr8s/tests/test_objects.py
index 0f3c3a2..25dafcd 100644
--- a/kr8s/tests/test_objects.py
+++ b/kr8s/tests/test_objects.py
@@ -15,6 +15,7 @@
 
 import kr8s
 from kr8s._async_utils import anext
+from kr8s._exceptions import NotFoundError
 from kr8s._exec import CompletedExec, ExecError
 from kr8s.asyncio.objects import (
     APIObject,
@@ -658,6 +659,29 @@ async def test_node():
         await node.uncordon()
 
 
+async def test_node_taint():
+    api = await kr8s.asyncio.api()
+    nodes = [node async for node in api.get("nodes")]
+    assert len(nodes) > 0
+    node = nodes[0]
+
+    # Remove existing taints just in case they still exist
+    for taint in node.taints:
+        await node.taint(key=taint["key"], value=taint["value"], effect="NoSchedule-")
+    assert not node.taints
+
+    await node.taint(key="key1", value="value1", effect="NoSchedule")
+    await node.taint(key="key2", value="value2", effect="NoSchedule")
+    assert len(node.taints) == 2
+
+    await node.taint(key="key1", value="value1", effect="NoSchedule-")
+    await node.taint(key="key2", value="value2", effect="NoSchedule-")
+    assert not node.taints
+
+    with pytest.raises(NotFoundError):
+        await node.taint(key="key123", value="value1", effect="NoSchedule-")
+
+
 async def test_service_proxy():
     api = await kr8s.asyncio.api()
     service = await anext(api.get("services", "kubernetes"))