From fd3f533dfc2d4a67a8679dd2666c249f0869f500 Mon Sep 17 00:00:00 2001
From: Eran Kampf <205185+ekampf@users.noreply.github.com>
Date: Fri, 1 Mar 2024 08:43:48 -0800
Subject: [PATCH] fix: Improved `twingate_connector_recreate_pod` logic (#179)

---
 app/handlers/handlers_connectors.py           | 45 ++++++++-----------
 app/handlers/tests/test_handlers_connector.py |  7 ++-
 tests_integration/test_connector_flows.py     |  6 ++-
 3 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/app/handlers/handlers_connectors.py b/app/handlers/handlers_connectors.py
index 08dc6a72..da0d68da 100644
--- a/app/handlers/handlers_connectors.py
+++ b/app/handlers/handlers_connectors.py
@@ -1,5 +1,3 @@
-import time
-
 import kopf
 import kubernetes.client
 import pendulum
@@ -85,17 +83,24 @@ def get_connector_secret(
     )
 
 
-def check_pod_exists(namespace: str, name: str) -> bool:
+def get_existing_pod(
+    namespace: str, name: str, kapi: kubernetes.client.CoreV1Api | None = None
+) -> kubernetes.client.V1Pod | None:
     try:
-        kapi = kubernetes.client.CoreV1Api()
-        kapi.read_namespaced_pod(name=name, namespace=namespace)
-        return True
+        kapi = kapi or kubernetes.client.CoreV1Api()
+        return kapi.read_namespaced_pod(name=name, namespace=namespace)
     except kubernetes.client.exceptions.ApiException as ex:
         if ex.status == 404:
-            return False
+            return None
         raise
 
 
+def check_pod_exists(
+    namespace: str, name: str, kapi: kubernetes.client.CoreV1Api | None = None
+) -> bool:
+    return bool(get_existing_pod(namespace, name, kapi=kapi))
+
+
 @kopf.on.create("twingateconnector")
 def twingate_connector_create(body, memo, logger, namespace, patch, **_):
     settings = memo.twingate_settings
@@ -265,27 +270,13 @@ def is_conflict_already_exists(apiex):
     kopf.adopt(pod, owner=body, strict=True, forced=True)
     kopf.label(pod, {"twingate.com/connector": crd.metadata.name})
 
-    retry_count = 0
     kapi = kubernetes.client.CoreV1Api()
-    while True:
-        try:
-            kapi.create_namespaced_pod(namespace=namespace, body=pod)
-            break
-        except kubernetes.client.exceptions.ApiException as apiex:  # noqa: PERF203
-            if is_conflict_already_exists(apiex):
-                # Pod is not deleted yet...  retry
-                logger.info(
-                    "Pod %s not deleted yet. Retrying (%s)...",
-                    pod.metadata.name,
-                    retry_count,
-                )
-                retry_count += 1
-                if retry_count > 3:
-                    raise
-                time.sleep(2)
-            else:
-                raise
+    if check_pod_exists(namespace, crd.metadata.name, kapi=kapi):
+        raise kopf.TemporaryError(
+            f"Pod {crd.metadata.name} not deleted yet. Retrying (%s)...", delay=1
+        )
 
+    kapi.create_namespaced_pod(namespace=namespace, body=pod)
     patch.meta["labels"] = {LABEL_CONNECTOR_POD_DELETED: None}
 
 
@@ -338,5 +329,7 @@ def twingate_connector_pod_deleted(body, spec, meta, logger, namespace, memo, **
     except kubernetes.client.exceptions.ApiException:
         logger.exception("Failed to annotate connector %s", owner["name"])
 
+    return success(msg="deleted")
+
 
 # endregion
diff --git a/app/handlers/tests/test_handlers_connector.py b/app/handlers/tests/test_handlers_connector.py
index 6a1ad5ec..0c9a4c40 100644
--- a/app/handlers/tests/test_handlers_connector.py
+++ b/app/handlers/tests/test_handlers_connector.py
@@ -319,8 +319,13 @@ def test_timer_check_image_version_with_imagepolicy_do_nothing_if_check_not_due(
     assert run.patch_mock.meta == {}
 
 
-def test_twingate_connector_recreate_pod(get_connector_and_crd, kopf_handler_runner):
+def test_twingate_connector_recreate_pod(
+    k8s_client_mock, get_connector_and_crd, kopf_handler_runner
+):
     connector, crd = get_connector_and_crd()
+
+    k8s_client_mock.read_namespaced_pod.return_value = None
+
     run = kopf_handler_runner(twingate_connector_recreate_pod, crd, MagicMock())
 
     run.kopf_adopt_mock.assert_called_once_with(
diff --git a/tests_integration/test_connector_flows.py b/tests_integration/test_connector_flows.py
index 82807f87..e666b4a5 100644
--- a/tests_integration/test_connector_flows.py
+++ b/tests_integration/test_connector_flows.py
@@ -89,7 +89,7 @@ def test_connector_flows(kopf_settings, kopf_runner_args, ci_run_number):
     assert runner.exit_code == 0
 
 
-def test_connector_flowes_image_change(kopf_settings, kopf_runner_args, ci_run_number):
+def test_connector_flows_image_change(kopf_settings, kopf_runner_args, ci_run_number):
     connector_name = f"test-connector-image-{ci_run_number}"
     OBJ = f"""
         apiVersion: twingate.com/v1beta
@@ -151,7 +151,9 @@ def test_connector_flowes_image_change(kopf_settings, kopf_runner_args, ci_run_n
     assert runner.exit_code == 0
 
 
-def test_pod_gone_while_operator_down(kopf_settings, kopf_runner_args, ci_run_number):
+def test_connector_flows_pod_gone_while_operator_down(
+    kopf_settings, kopf_runner_args, ci_run_number
+):
     connector_name = f"test-connector-gone-{ci_run_number}"
     OBJ = f"""
         apiVersion: twingate.com/v1beta