From 3bd0a0f644337d282efd432ce6112ef9488d31fd Mon Sep 17 00:00:00 2001
From: Doruk Ozturk <dozturk@amazon.com>
Date: Wed, 15 Feb 2023 12:09:01 -0500
Subject: [PATCH] refactor: Simplify tests

---
 .../cluster_wide/scalability/control_plane.py | 39 +++++++------------
 tests/test_scalability_control_plane.py       | 32 +++++++--------
 2 files changed, 31 insertions(+), 40 deletions(-)

diff --git a/hardeneks/cluster_wide/scalability/control_plane.py b/hardeneks/cluster_wide/scalability/control_plane.py
index a96d1a6..7cae80a 100644
--- a/hardeneks/cluster_wide/scalability/control_plane.py
+++ b/hardeneks/cluster_wide/scalability/control_plane.py
@@ -1,33 +1,24 @@
-from ...resources import Resources
-from rich.console import Console
+import re
 from rich.panel import Panel
-from rich import print
 import kubernetes
 
-console = Console()
+from hardeneks import console
+from ...resources import Resources
 
-def _get_version() -> str:
-    client = kubernetes.client.VersionApi()
-    version = client.get_code()
-    return version
 
 def check_EKS_version(resources: Resources):
-    version = _get_version()
-    major = int(version.major)
+    client = kubernetes.client.VersionApi()
+    version = client.get_code()
     minor = version.minor
-    last_char = version.minor[-1]
-    if last_char == "+":
-        minor = int(version.minor[:-1])
-    else:
-        minor = int(minor)
-
-    good = False
 
-    if major >= 1 and minor >= 24:
-        good = True
-
-    if good == False:
-        print(Panel("[red] Current Version == " + version.major + "." + version.minor + "", title="EKS Version Should be greater or equal too 1.24"))
+    if int(re.sub("[^0-9]", "", minor)) < 24:
+        console.print(
+            Panel(
+                f"[red]EKS Version Should be greater or equal too 1.24. Current Version == {version.major}.{version.minor}",
+                subtitle="[link=https://aws.github.io/aws-eks-best-practices/scalability/docs/control-plane/#use-eks-124-or-above]Click to see the guide[/link]",
+            )
+        )
         console.print()
-    
-    return good
\ No newline at end of file
+        return False
+
+    return True
diff --git a/tests/test_scalability_control_plane.py b/tests/test_scalability_control_plane.py
index fda3731..13ffaa0 100644
--- a/tests/test_scalability_control_plane.py
+++ b/tests/test_scalability_control_plane.py
@@ -1,23 +1,23 @@
-import pytest
 from hardeneks.resources import Resources
-from hardeneks.resources import Map
 from unittest.mock import patch
- 
-from hardeneks.cluster_wide.scalability.control_plane import (
-    check_EKS_version
-)
 
-@patch("hardeneks.cluster_wide.scalability.control_plane._get_version")
+from hardeneks.cluster_wide.scalability.control_plane import check_EKS_version
+
+
+class Version:
+    def __init__(self, minor):
+        self.major = 1
+        self.minor = minor
+
+
+@patch("kubernetes.client.VersionApi.get_code")
 def test_check_EKS_version(mocked_client):
     namespaced_resources = Resources(
         "some_region", "some_context", "some_cluster", []
     )
-    mocked_client.return_value = Map({'major': '1', 'minor': "23+"})
-    good = check_EKS_version(namespaced_resources)
-    assert good == False, "Value should be False"
-    mocked_client.return_value = Map({'major': '1', 'minor': "24+"})
-    good = check_EKS_version(namespaced_resources)
-    assert good == True, "Value should be True"
-    mocked_client.return_value = Map({'major': '1', 'minor': "24"})
-    good = check_EKS_version(namespaced_resources)
-    assert good == True, "Value should be True and Handle without +"
\ No newline at end of file
+    mocked_client.return_value = Version("23+")
+    assert not check_EKS_version(namespaced_resources)
+    mocked_client.return_value = Version("24+")
+    assert check_EKS_version(namespaced_resources)
+    mocked_client.return_value = Version("24")
+    assert check_EKS_version(namespaced_resources)