From 18d38aac7164c323ad025f63bb5e16fd729020d1 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Sat, 7 Dec 2024 10:03:52 -0500 Subject: [PATCH 1/5] Add taint and taints to Node object --- kr8s/_objects.py | 31 ++++++++++++++++++++++++++++++- kr8s/tests/test_objects.py | 15 +++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/kr8s/_objects.py b/kr8s/_objects.py index f8700d9..d04cb11 100644 --- a/kr8s/_objects.py +++ b/kr8s/_objects.py @@ -378,7 +378,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": @@ -824,6 +826,33 @@ 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.""" + if effect.endswith("-"): + effect = effect[:-1] + await self.async_patch( + [ + { + "op": "remove", + "path": "/spec/taints", + "value": [{"key": key, "value": value, "effect": effect}], + } + ], + type="json", + ) + else: + await self.async_patch( + {"spec": {"taints": [{"effect": effect, "key": key, "value": value}]}} + ) + + @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 be51261..750ec97 100644 --- a/kr8s/tests/test_objects.py +++ b/kr8s/tests/test_objects.py @@ -648,6 +648,21 @@ 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] + + await node.taint(key="key1", value="value1", effect="NoSchedule") + assert any( + taint["key"] == "key1" and taint["value"] == "value1" for taint in node.taints + ) + + await node.taint(key="key1", value="value1", effect="NoSchedule-") + assert not node.taints + + async def test_service_proxy(): api = await kr8s.asyncio.api() [service] = await api.get("services", "kubernetes") From 1bcdefab6bb3bcf8defc1ddfd0d9283b681482c3 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Mon, 16 Dec 2024 20:46:29 -0500 Subject: [PATCH 2/5] Use simplier implementation --- kr8s/_objects.py | 23 ++++++++++------------- kr8s/tests/test_objects.py | 12 ++++++++---- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/kr8s/_objects.py b/kr8s/_objects.py index d04cb11..82f31e5 100644 --- a/kr8s/_objects.py +++ b/kr8s/_objects.py @@ -829,21 +829,18 @@ async def uncordon(self) -> None: async def taint(self, key: str, value: str, *, effect: str) -> None: """Taint a node.""" if effect.endswith("-"): + # Remove taint with key effect = effect[:-1] - await self.async_patch( - [ - { - "op": "remove", - "path": "/spec/taints", - "value": [{"key": key, "value": value, "effect": effect}], - } - ], - type="json", - ) + 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: - await self.async_patch( - {"spec": {"taints": [{"effect": effect, "key": key, "value": value}]}} - ) + taints = list(self.taints) + [ + {"key": key, "value": value, "effect": effect} + ] + + await self.async_patch({"spec": {"taints": taints}}) @property def taints(self) -> Box: diff --git a/kr8s/tests/test_objects.py b/kr8s/tests/test_objects.py index 750ec97..299ed8d 100644 --- a/kr8s/tests/test_objects.py +++ b/kr8s/tests/test_objects.py @@ -14,6 +14,7 @@ import pytest import kr8s +from kr8s._exceptions import NotFoundError from kr8s._exec import CompletedExec, ExecError from kr8s.asyncio.objects import ( APIObject, @@ -650,18 +651,21 @@ async def test_node(): async def test_node_taint(): api = await kr8s.asyncio.api() - nodes = [node async for node in api.get("nodes")] + nodes = await api.get("nodes") assert len(nodes) > 0 node = nodes[0] await node.taint(key="key1", value="value1", effect="NoSchedule") - assert any( - taint["key"] == "key1" and taint["value"] == "value1" for taint in node.taints - ) + 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() From 2ba0f2d3204222a89ec505613e69583cb8b4756d Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Mon, 16 Dec 2024 20:54:00 -0500 Subject: [PATCH 3/5] Fix taints --- kr8s/tests/test_objects.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kr8s/tests/test_objects.py b/kr8s/tests/test_objects.py index cf40454..6e183e2 100644 --- a/kr8s/tests/test_objects.py +++ b/kr8s/tests/test_objects.py @@ -661,7 +661,7 @@ async def test_node(): async def test_node_taint(): api = await kr8s.asyncio.api() - nodes = await api.get("nodes") + nodes = [node async for node in api.get("nodes")] assert len(nodes) > 0 node = nodes[0] From 6f4886d2cd1f3d99237ce5b255f8121fad8bc3a3 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Mon, 16 Dec 2024 20:56:22 -0500 Subject: [PATCH 4/5] Update tests --- kr8s/tests/test_objects.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/kr8s/tests/test_objects.py b/kr8s/tests/test_objects.py index 6e183e2..25dafcd 100644 --- a/kr8s/tests/test_objects.py +++ b/kr8s/tests/test_objects.py @@ -665,6 +665,11 @@ async def test_node_taint(): 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 From e128abba2835dbd3b98f4227ee719a90d887fe12 Mon Sep 17 00:00:00 2001 From: Jacob Tomlinson Date: Tue, 17 Dec 2024 11:09:51 +0000 Subject: [PATCH 5/5] Refresh state before modifying taints --- kr8s/_objects.py | 1 + 1 file changed, 1 insertion(+) diff --git a/kr8s/_objects.py b/kr8s/_objects.py index cc1bca3..d7ccbed 100644 --- a/kr8s/_objects.py +++ b/kr8s/_objects.py @@ -832,6 +832,7 @@ async def uncordon(self) -> None: 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]