fix: Adjust node_type check for get ami (#101)

eksupgrade/__init__.py

@@ -5,4 +5,4 @@
 
 """
 
-__version__: str = "0.8.1"
+__version__: str = "0.8.2"

eksupgrade/src/eks_get_image_type.py

@@ -20,14 +20,12 @@ def image_type(node_type: str, image_id: str, region: str) -> Optional[str]:
         {"Name": "is-public", "Values": ["true"]},
     ]
 
-    if node_type == "amazon linux 2":
-        filters.append({"Name": "name", "Values": ["amazon-eks-node-*"]})
-    elif node_type == "ubuntu":
-        filters.append({"Name": "name", "Values": ["ubuntu-eks/k8s_*"]})
-    elif node_type == "bottlerocket":
+    if "amazon linux 2" in node_type:
+        filters.append({"Name": "name", "Values": ["amazon-eks-node*"]})
+    elif "bottlerocket" in node_type:
         filters.append({"Name": "name", "Values": ["bottlerocket-aws-k8s-*"]})
-    elif node_type == "windows":
-        filters.append({"Name": "name", "Values": ["Windows_Server-*-English-*-EKS_Optimized-*"]})
+    elif "windows" in node_type:
+        filters.append({"Name": "name", "Values": ["Windows_Server-*-English-*-EKS_Optimized*"]})
     else:
         echo_warning(f"Node type: {node_type} is unsupported  - Image ID: {image_id}")
         return None

eksupgrade/src/preflight_module.py

@@ -834,8 +834,6 @@ def iscustomami(node_type, Presentversion, image_id, region):
 
     if node_type == "Amazon Linux 2":
         filters.append({"Name": "name", "Values": [f"amazon-eks-node-{Presentversion}*"]})
-    elif "ubuntu" in node_type.lower():
-        filters.append({"Name": "name", "Values": [f"ubuntu-eks/k8s_{Presentversion}*"]})
     elif "bottlerocket" in node_type.lower():
         filters.append({"Name": "name", "Values": [f"bottlerocket-aws-k8s-{Presentversion}*"]})
     elif "windows" in node_type.lower():

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "eksupgrade"
-version = "0.8.1"
+version = "0.8.2"
 description = "The Amazon EKS cluster upgrade utility"
 authors = ["EKS Upgrade Maintainers <[email protected]>"]
 readme = "README.md"

tests/conftest.py

@@ -4,7 +4,7 @@
 
 import boto3
 import pytest
-from moto import mock_eks, mock_sts
+from moto import mock_ec2, mock_eks, mock_sts
 
 
 @pytest.fixture
@@ -30,6 +30,14 @@ def sts_client(aws_creds, region) -> Generator[Any, None, None]:
         yield client
 
 
+@pytest.fixture
+def ec2_client(aws_creds, region) -> Generator[Any, None, None]:
+    """Mock the EKS boto client."""
+    with mock_ec2():
+        client = boto3.client("ec2", region_name=region)
+        yield client
+
+
 @pytest.fixture
 def eks_client(aws_creds, region) -> Generator[Any, None, None]:
     """Mock the EKS boto client."""

tests/test_eks_get_image_type.py

@@ -0,0 +1,20 @@
+"""Test EKS Upgrade get image type specific logic."""
+from typing import Optional
+
+import pytest
+
+from eksupgrade.src.eks_get_image_type import image_type
+
+
+@pytest.mark.parametrize(
+    "node_type,image_id",
+    [
+        ("windows server 2019 datacenter ", "ami-ekswin"),
+        ("windows server 2022", "ami-ekswin"),
+        ("amazon linux 2", "ami-ekslinux"),
+    ],
+)
+def test_image_type(ec2_client, region, node_type, image_id) -> None:
+    """Test the image_type method."""
+    ami_id: Optional[str] = image_type(node_type=node_type, image_id=image_id, region=region)
+    assert ami_id