Merge pull request #88 from samuel-esp/fix-scaledobject-keyerror

Fix KEDA ScaledObject KeyError - Scale Up
caas-team · Aug 22, 2024 · b299be8 · b299be8
2 parents 8def356 + a6bd26b
commit b299be8
Show file tree

Hide file tree

Showing 2 changed files with 167 additions and 5 deletions.
diff --git a/kube_downscaler/scaler.py b/kube_downscaler/scaler.py
@@ -589,10 +589,11 @@ def scale_up(
             f"Scaling up {resource.kind} {resource.namespace}/{resource.name} from {replicas} to {original_replicas} replicas (uptime: {uptime}, downtime: {downtime})"
         )
     elif resource.kind == "ScaledObject":
-        if resource.annotations[ScaledObject.last_keda_pause_annotation_if_present] is not None:
-            paused_replicas = resource.annotations[ScaledObject.last_keda_pause_annotation_if_present]
-            resource.annotations[ScaledObject.keda_pause_annotation] = paused_replicas
-            resource.annotations[ScaledObject.last_keda_pause_annotation_if_present] = None
+        if ScaledObject.keda_pause_annotation in resource.annotations:
+            if resource.annotations[ScaledObject.keda_pause_annotation] is not None:
+                paused_replicas = resource.annotations[ScaledObject.last_keda_pause_annotation_if_present]
+                resource.annotations[ScaledObject.keda_pause_annotation] = paused_replicas
+                resource.annotations[ScaledObject.last_keda_pause_annotation_if_present] = None
         else:
             resource.annotations[ScaledObject.keda_pause_annotation] = None
         logger.info(

diff --git a/tests/test_autoscale_resource.py b/tests/test_autoscale_resource.py
@@ -11,6 +11,7 @@
 from pykube import HorizontalPodAutoscaler
 
 from kube_downscaler.resources.stack import Stack
+from kube_downscaler.resources.keda import ScaledObject
 from kube_downscaler.scaler import autoscale_resource
 from kube_downscaler.scaler import DOWNSCALE_PERIOD_ANNOTATION
 from kube_downscaler.scaler import DOWNTIME_REPLICAS_ANNOTATION
@@ -1143,5 +1144,165 @@ def test_upscale_daemonset_with_autoscaling():
         matching_labels=frozenset([re.compile("")]),
     )
 
-    print(ds.obj)
     assert ds.obj["spec"]["template"]["spec"]["nodeSelector"]["kube-downscaler-non-existent"] == None
+
+
+def test_downscale_scaledobject_with_pause_annotation_already_present():
+    # Create a ScaledObject with the annotation present
+    so = ScaledObject(
+        None,
+        {
+            "metadata": {
+                "name": "scaledobject-1",
+                "namespace": "default",
+                "creationTimestamp": "2023-08-21T10:00:00Z",
+                "annotations": {
+                    "autoscaling.keda.sh/paused-replicas": "3"
+                }
+            },
+            "spec": {}
+        }
+    )
+
+    now = datetime.strptime("2023-08-21T10:30:00Z", "%Y-%m-%dT%H:%M:%SZ").replace(
+        tzinfo=timezone.utc
+    )
+
+
+    autoscale_resource(
+        so,
+        upscale_target_only=False,
+        upscale_period="never",
+        downscale_period="never",
+        default_uptime="never",
+        default_downtime="always",
+        forced_uptime=False,
+        forced_downtime=False,
+        dry_run=True,
+        now=now,
+        matching_labels=frozenset([re.compile("")]),
+    )
+
+    # Check if the annotations have been correctly updated
+    assert so.annotations[ScaledObject.keda_pause_annotation] == "0"
+    assert so.annotations[ScaledObject.last_keda_pause_annotation_if_present] == "3"
+
+
+def test_upscale_scaledobject_with_pause_annotation_already_present():
+    so = ScaledObject(
+        None,
+        {
+            "metadata": {
+                "name": "scaledobject-1",
+                "namespace": "default",
+                "creationTimestamp": "2023-08-21T10:00:00Z",
+                "annotations": {
+                    "autoscaling.keda.sh/paused-replicas": "0",  # Paused replicas
+                    "downscaler/original-pause-replicas": "3",  # Original replicas before pause
+                    "downscaler/original-replicas": "3",  # Keeping track of original replicas
+                }
+            },
+            "spec": {}
+        }
+    )
+
+
+    now = datetime.strptime("2023-08-21T10:30:00Z", "%Y-%m-%dT%H:%M:%SZ").replace(
+        tzinfo=timezone.utc
+    )
+
+
+    autoscale_resource(
+        so,
+        upscale_target_only=False,
+        upscale_period="never",
+        downscale_period="never",
+        default_uptime="always",
+        default_downtime="never",
+        forced_uptime=False,
+        forced_downtime=False,
+        dry_run=True,
+        now=now,
+        matching_labels=frozenset([re.compile("")]),
+    )
+
+    # Check if the annotations have been correctly updated for the upscale operation
+    assert so.annotations[ScaledObject.keda_pause_annotation] == "3"
+    assert so.annotations.get(ScaledObject.last_keda_pause_annotation_if_present) is None
+
+def test_downscale_scaledobject_without_pause_annotation():
+    so = ScaledObject(
+        None,
+        {
+            "metadata": {
+                "name": "scaledobject-1",
+                "namespace": "default",
+                "creationTimestamp": "2023-08-21T10:00:00Z",
+            },
+            "spec": {}
+        }
+    )
+
+    now = datetime.strptime("2023-08-21T10:30:00Z", "%Y-%m-%dT%H:%M:%SZ").replace(
+        tzinfo=timezone.utc
+    )
+
+    autoscale_resource(
+        so,
+        upscale_target_only=False,
+        upscale_period="never",
+        downscale_period="never",
+        default_uptime="never",
+        default_downtime="always",
+        forced_uptime=False,
+        forced_downtime=False,
+        dry_run=True,
+        now=now,
+        matching_labels=frozenset([re.compile("")]),
+    )
+
+    # Check if the annotations have been correctly updated
+    assert so.annotations[ScaledObject.keda_pause_annotation] == "0"
+
+
+def test_upscale_scaledobject_without_pause_annotation():
+    so = ScaledObject(
+        None,
+        {
+            "metadata": {
+                "name": "scaledobject-1",
+                "namespace": "default",
+                "creationTimestamp": "2023-08-21T10:00:00Z",
+                "annotations": {
+                    "autoscaling.keda.sh/paused-replicas": "0",
+                    "downscaler/original-pause-replicas": "3",
+                    "downscaler/original-replicas": "3",
+                }
+            },
+            "spec": {}
+        }
+    )
+
+
+    now = datetime.strptime("2023-08-21T10:30:00Z", "%Y-%m-%dT%H:%M:%SZ").replace(
+        tzinfo=timezone.utc
+    )
+
+
+    autoscale_resource(
+        so,
+        upscale_target_only=False,
+        upscale_period="never",
+        downscale_period="never",
+        default_uptime="always",
+        default_downtime="never",
+        forced_uptime=False,
+        forced_downtime=False,
+        dry_run=True,
+        now=now,
+        matching_labels=frozenset([re.compile("")]),
+    )
+
+    # Check if the annotations have been correctly updated for the upscale operation
+    assert so.annotations[ScaledObject.keda_pause_annotation] == "3"
+    assert so.annotations.get(ScaledObject.last_keda_pause_annotation_if_present) is None