From 739f3f4091c5f9d7da4fc9d8f081a3e347d6d485 Mon Sep 17 00:00:00 2001
From: Bhavya Sharma <bhavyssh@amazon.com>
Date: Tue, 31 Mar 2026 09:11:53 -0700
Subject: [PATCH] feat(sagemaker): Add SageMakerConditionOperator and
 SageMakerFailOperator

---
 providers/amazon/docs/operators/sagemaker.rst |  32 +++
 .../amazon/aws/operators/sagemaker.py         | 266 +++++++++++++++++-
 .../amazon/aws/example_sagemaker_condition.py | 179 ++++++++++++
 .../aws/operators/test_sagemaker_condition.py | 222 +++++++++++++++
 4 files changed, 698 insertions(+), 1 deletion(-)
 create mode 100644 providers/amazon/tests/system/amazon/aws/example_sagemaker_condition.py
 create mode 100644 providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_condition.py

diff --git a/providers/amazon/docs/operators/sagemaker.rst b/providers/amazon/docs/operators/sagemaker.rst
index 5103b868455bb..71f61a3289098 100644
--- a/providers/amazon/docs/operators/sagemaker.rst
+++ b/providers/amazon/docs/operators/sagemaker.rst
@@ -385,6 +385,38 @@ you can use :class:`~airflow.providers.amazon.aws.sensors.sagemaker.SageMakerPro
     :start-after: [START howto_sensor_sagemaker_processing]
     :end-before: [END howto_sensor_sagemaker_processing]
 
+.. _howto/operator:SageMakerConditionOperator:
+
+Branch a DAG based on condition evaluation
+==========================================
+
+To branch an Airflow DAG based on upstream task outputs you can use
+:class:`~airflow.providers.amazon.aws.operators.sagemaker.SageMakerConditionOperator`.
+
+Simple usage with flat parameters (single condition):
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_sagemaker_condition.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_sagemaker_condition_flat]
+    :end-before: [END howto_operator_sagemaker_condition_flat]
+
+Advanced usage with conditions list (multiple AND-ed conditions):
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_sagemaker_condition.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_sagemaker_condition]
+    :end-before: [END howto_operator_sagemaker_condition]
+
+Using Not and Or conditions:
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_sagemaker_condition.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_sagemaker_condition_not_or]
+    :end-before: [END howto_operator_sagemaker_condition_not_or]
+
 Reference
 ---------
 
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
index ff5e754913a9d..fe1597f5b1f07 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
@@ -41,7 +41,12 @@
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.sagemaker import ApprovalStatus
 from airflow.providers.amazon.aws.utils.tags import format_tags
-from airflow.providers.common.compat.sdk import AirflowException, conf
+from airflow.providers.common.compat.sdk import (
+    AirflowException,
+    AirflowFailException,
+    BaseBranchOperator,
+    conf,
+)
 from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:
@@ -1991,3 +1996,262 @@ def execute(self, context):
             self.hook.conn.get_waiter("notebook_instance_in_service").wait(
                 NotebookInstanceName=self.instance_name
             )
+
+
+class SageMakerConditionOperator(BaseBranchOperator):
+    """
+    Evaluates a single condition or a list of conditions, and routes tasks based on the result.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:SageMakerConditionOperator`
+
+    :param condition_type: Condition type for the simple (flat) interface.
+        Valid types: ``Equals``, ``GreaterThan``, ``GreaterThanOrEqualTo``,
+        ``LessThan``, ``LessThanOrEqualTo``, ``In``.
+        Mutually exclusive with ``conditions``.
+    :param left_value: Left operand for the flat interface. For ``In`` conditions
+        this is the value to check membership of.
+    :param right_value: Right operand for the flat interface. For ``In`` conditions
+        this is the list of allowed values.
+    :param conditions: List of condition dicts to evaluate (AND-ed together).
+        Each dict must have a ``type`` key. Must not be empty.
+        Mutually exclusive with ``condition_type``/``left_value``/``right_value``.
+    :param if_task_ids: Task ID(s) to execute when all conditions are True.
+    :param else_task_ids: Task ID(s) to execute when any condition is False.
+        If omitted, the task fails with ``AirflowFailException`` when conditions are not met.
+    """
+
+    _VALID_FLAT_TYPES: ClassVar[set[str]] = {
+        "Equals",
+        "GreaterThan",
+        "GreaterThanOrEqualTo",
+        "LessThan",
+        "LessThanOrEqualTo",
+        "In",
+    }
+
+    template_fields: Sequence[str] = (
+        "condition_type",
+        "left_value",
+        "right_value",
+        "conditions",
+        "if_task_ids",
+        "else_task_ids",
+    )
+
+    def __init__(
+        self,
+        *,
+        condition_type: str | None = None,
+        left_value: Any = None,
+        right_value: Any = None,
+        conditions: list[dict] | None = None,
+        if_task_ids: str | list[str],
+        else_task_ids: str | list[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        has_flat = condition_type is not None
+        has_list = conditions is not None
+
+        if has_flat and has_list:
+            raise ValueError(
+                "Cannot use 'condition_type' and 'conditions' together. "
+                "Use 'condition_type' with 'left_value'/'right_value' for a single condition, "
+                "or 'conditions' for multiple/nested conditions."
+            )
+        if not has_flat and not has_list:
+            raise ValueError(
+                "Missing condition: provide 'condition_type' with 'left_value'/'right_value' "
+                "for a single condition, or 'conditions' for multiple/nested conditions."
+            )
+
+        if has_flat:
+            if condition_type not in self._VALID_FLAT_TYPES:
+                raise ValueError(
+                    f"Unknown condition_type '{condition_type}'. "
+                    f"Expected one of: {', '.join(sorted(self._VALID_FLAT_TYPES))}."
+                )
+            self.condition_type: str | None = condition_type
+            self.left_value = left_value
+            self.right_value = right_value
+            if condition_type == "In":
+                self.conditions: list[dict[str, Any]] = [
+                    {"type": "In", "value": left_value, "in_values": right_value}
+                ]
+            else:
+                self.conditions = [
+                    {"type": condition_type, "left_value": left_value, "right_value": right_value}
+                ]
+        else:
+            self.condition_type = None
+            self.left_value = None
+            self.right_value = None
+            self.conditions = conditions  # type: ignore[assignment]
+
+        if not self.conditions:
+            raise ValueError("At least 1 condition is required, but got an empty list.")
+        self.if_task_ids = [if_task_ids] if isinstance(if_task_ids, str) else if_task_ids
+        self.else_task_ids = [else_task_ids] if isinstance(else_task_ids, str) else (else_task_ids or [])
+
+    @staticmethod
+    def _cast(value: Any) -> Any:
+        """
+        Cast Jinja-rendered string values to appropriate Python types.
+
+        This is a compatibility shim for environments where
+        ``render_template_as_native_obj=True`` is not available at the DAG or
+        task level (e.g., YAML DAGs). Once task-level native rendering
+        is widely supported, this method can be removed in favor of letting
+        Airflow handle the casting natively.
+
+        - Numeric strings become int or float.
+        - ``"true"``/``"false"`` become booleans.
+        - ``"None"`` becomes ``None`` (common when ``xcom_pull`` returns nothing).
+        - Other strings are returned unchanged.
+        - Non-string types pass through as-is.
+        """
+        if not isinstance(value, str):
+            return value
+        if value == "None":
+            return None
+        try:
+            return int(value)
+        except (ValueError, TypeError):
+            pass
+        try:
+            return float(value)
+        except (ValueError, TypeError):
+            pass
+        if value.lower() == "true":
+            return True
+        if value.lower() == "false":
+            return False
+        return value
+
+    _COMPARISON_OPERATORS: ClassVar[dict[str, Callable[[Any, Any], bool]]] = {
+        "Equals": lambda left, right: left == right,
+        "GreaterThan": lambda left, right: left > right,
+        "GreaterThanOrEqualTo": lambda left, right: left >= right,
+        "LessThan": lambda left, right: left < right,
+        "LessThanOrEqualTo": lambda left, right: left <= right,
+    }
+
+    def _evaluate(self, condition: dict, depth: int = 0) -> bool:
+        """
+        Recursively evaluate a single condition dict.
+
+        :param condition: A condition dictionary with a ``type`` key.
+        :param depth: Current nesting depth (used for log indentation only).
+        :returns: Boolean result of the condition evaluation.
+        """
+        log_indent = "  " * depth
+        try:
+            condition_type = condition["type"]
+        except KeyError:
+            raise ValueError("Condition dict is missing required key 'type'.")
+
+        if condition_type in self._COMPARISON_OPERATORS:
+            try:
+                left = self._cast(condition["left_value"])
+                right = self._cast(condition["right_value"])
+            except KeyError as e:
+                raise ValueError(f"Condition '{condition_type}' missing required key {e}.") from None
+
+            # None check — likely an XCom that was not pushed
+            if left is None or right is None:
+                raise TypeError(
+                    f"Condition '{condition_type}' received None: left={left!r}, right={right!r}. "
+                    "This usually means the upstream task did not run or did not push a value to XCom."
+                )
+
+            # Type compatibility check
+            left_type = type(left)
+            right_type = type(right)
+            numeric_types = (int, float)
+            left_is_numeric = isinstance(left, numeric_types) and not isinstance(left, bool)
+            right_is_numeric = isinstance(right, numeric_types) and not isinstance(right, bool)
+
+            if not (left_is_numeric and right_is_numeric) and left_type is not right_type:
+                raise TypeError(
+                    f"Cannot compare {left_type.__name__} ({left!r}) with {right_type.__name__} ({right!r}) "
+                    f"in condition '{condition_type}'. Both values must be the same type."
+                )
+
+            comparison_result = self._COMPARISON_OPERATORS[condition_type](left, right)
+            self.log.info(
+                "%s%s: %r %s %r -> %s",
+                log_indent,
+                condition_type,
+                left,
+                condition_type,
+                right,
+                comparison_result,
+            )
+            return comparison_result
+
+        if condition_type == "In":
+            try:
+                query_value = self._cast(condition["value"])
+                allowed_values = [self._cast(val) for val in condition["in_values"]]
+            except KeyError as e:
+                raise ValueError(f"Condition '{condition_type}' missing required key {e}.") from None
+            membership_result = query_value in allowed_values
+            self.log.info("%sIn: %r in %r -> %s", log_indent, query_value, allowed_values, membership_result)
+            return membership_result
+
+        if condition_type == "Not":
+            try:
+                inner_condition = condition["condition"]
+            except KeyError as e:
+                raise ValueError(f"Condition '{condition_type}' missing required key {e}.") from None
+            inner_result = self._evaluate(inner_condition, depth + 1)
+            negated_result = not inner_result
+            self.log.info("%sNot: not %s -> %s", log_indent, inner_result, negated_result)
+            return negated_result
+
+        if condition_type == "Or":
+            try:
+                inner_conditions = condition["conditions"]
+            except KeyError as e:
+                raise ValueError(f"Condition '{condition_type}' missing required key {e}.") from None
+            inner_results = [self._evaluate(inner_cond, depth + 1) for inner_cond in inner_conditions]
+            or_result = any(inner_results)
+            self.log.info("%sOr: any(%r) -> %s", log_indent, inner_results, or_result)
+            return or_result
+
+        raise ValueError(f"Unknown condition type '{condition_type}'.")
+
+    def choose_branch(self, context: Context) -> list[str]:
+        """
+        Evaluate all conditions and return the appropriate branch task IDs.
+
+        :param context: Airflow context dictionary.
+        :returns: ``if_task_ids`` when all conditions are True, ``else_task_ids`` otherwise.
+        """
+        condition_count = len(self.conditions)
+        self.log.info("Evaluating %d condition(s).", condition_count)
+
+        evaluation_results = [self._evaluate(condition) for condition in self.conditions]
+        all_conditions_met = all(evaluation_results)
+
+        if all_conditions_met:
+            self.log.info(
+                "All %d condition(s) evaluated to True. Routing to if_task_ids=%r.",
+                condition_count,
+                self.if_task_ids,
+            )
+            return self.if_task_ids
+
+        if not self.else_task_ids:
+            raise AirflowFailException(
+                f"Condition check failed in task '{self.task_id}': results={evaluation_results}"
+            )
+
+        self.log.info(
+            "Not all conditions are True (results=%r). Routing to else_task_ids=%r.",
+            evaluation_results,
+            self.else_task_ids,
+        )
+        return self.else_task_ids
diff --git a/providers/amazon/tests/system/amazon/aws/example_sagemaker_condition.py b/providers/amazon/tests/system/amazon/aws/example_sagemaker_condition.py
new file mode 100644
index 0000000000000..90a392fa0c1cb
--- /dev/null
+++ b/providers/amazon/tests/system/amazon/aws/example_sagemaker_condition.py
@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+System test for SageMakerConditionOperator.
+
+This operator evaluates conditions against XCom values passed from upstream tasks.
+
+The DAG simulates an ML accuracy-gate workflow:
+
+1. ``produce_metrics`` pushes a dict of metrics to XCom.
+2. ``check_accuracy`` uses SageMakerConditionOperator to branch:
+   - accuracy >= 0.9 AND loss < 0.1 -> ``deploy_model``
+   - otherwise -> ``retrain_model``
+3. Only the correct branch task runs; the other is skipped.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.sagemaker import SageMakerConditionOperator
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, chain, task
+else:
+    from airflow.decorators import task  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+
+from system.amazon.aws.utils import SystemTestContextBuilder
+
+DAG_ID = "example_sagemaker_condition"
+
+sys_test_context_task = SystemTestContextBuilder().build()
+
+
+with DAG(
+    DAG_ID,
+    schedule="@once",
+    start_date=datetime(2021, 1, 1),
+    catchup=False,
+) as dag:
+    test_context = sys_test_context_task()
+
+    # TEST SETUP: push simulated ML metrics to XCom
+
+    @task
+    def produce_metrics():
+        """Simulate an ML training job that returns accuracy and loss metrics."""
+        return {"accuracy": 0.95, "loss": 0.04}
+
+    metrics = produce_metrics()
+
+    # TEST BODY
+
+    # [START howto_operator_sagemaker_condition]
+    check_accuracy = SageMakerConditionOperator(
+        task_id="check_accuracy",
+        conditions=[
+            {
+                "type": "GreaterThanOrEqualTo",
+                "left_value": "{{ ti.xcom_pull(task_ids='produce_metrics')['accuracy'] }}",
+                "right_value": 0.9,
+            },
+            {
+                "type": "LessThan",
+                "left_value": "{{ ti.xcom_pull(task_ids='produce_metrics')['loss'] }}",
+                "right_value": 0.1,
+            },
+        ],
+        if_task_ids=["deploy_model"],
+        else_task_ids=["retrain_model"],
+    )
+    # [END howto_operator_sagemaker_condition]
+
+    @task
+    def deploy_model():
+        """Placeholder: model meets quality bar, proceed to deployment."""
+        return "deployed"
+
+    @task
+    def retrain_model():
+        """Placeholder: model does not meet quality bar, retrain."""
+        return "retrained"
+
+    # Scenario 2: condition evaluates to False -> else branch
+
+    @task
+    def produce_bad_metrics():
+        """Simulate a training job with poor accuracy."""
+        return {"accuracy": 0.5, "loss": 0.8}
+
+    bad_metrics = produce_bad_metrics()
+
+    # [START howto_operator_sagemaker_condition_flat]
+    check_bad_accuracy = SageMakerConditionOperator(
+        task_id="check_bad_accuracy",
+        condition_type="GreaterThanOrEqualTo",
+        left_value="{{ ti.xcom_pull(task_ids='produce_bad_metrics')['accuracy'] }}",
+        right_value=0.9,
+        if_task_ids=["should_not_run"],
+        else_task_ids=["should_run"],
+    )
+    # [END howto_operator_sagemaker_condition_flat]
+
+    @task
+    def should_not_run():
+        """This task should be skipped because accuracy < 0.9."""
+        return "error: should not have run"
+
+    @task
+    def should_run():
+        """This task should execute because accuracy < 0.9 -> else branch."""
+        return "correctly routed to else branch"
+
+    # Scenario 3: Or condition + Not condition
+
+    # [START howto_operator_sagemaker_condition_not_or]
+    check_logical = SageMakerConditionOperator(
+        task_id="check_logical",
+        conditions=[
+            {
+                "type": "Or",
+                "conditions": [
+                    {"type": "Equals", "left_value": 1, "right_value": 2},
+                    {"type": "Equals", "left_value": 3, "right_value": 3},
+                ],
+            },
+            {
+                "type": "Not",
+                "condition": {"type": "Equals", "left_value": "a", "right_value": "b"},
+            },
+        ],
+        if_task_ids=["logical_pass"],
+        else_task_ids=["logical_fail"],
+    )
+    # [END howto_operator_sagemaker_condition_not_or]
+
+    @task
+    def logical_pass():
+        """Or(1==2, 3==3) -> True AND Not(a==b) -> True -> if branch."""
+        return "logical conditions passed"
+
+    @task
+    def logical_fail():
+        return "error: logical conditions should have passed"
+
+    test_context >> [metrics, bad_metrics, check_logical]
+
+    chain(metrics, check_accuracy, [deploy_model(), retrain_model()])
+    chain(bad_metrics, check_bad_accuracy, [should_not_run(), should_run()])
+    chain(check_logical, [logical_pass(), logical_fail()])
+
+    from tests_common.test_utils.watcher import watcher
+
+    # This test needs watcher in order to properly mark success/failure
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
+
+from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_condition.py b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_condition.py
new file mode 100644
index 0000000000000..56d461e50a3fc
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_condition.py
@@ -0,0 +1,222 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from airflow.providers.amazon.aws.operators.sagemaker import SageMakerConditionOperator
+from airflow.providers.common.compat.sdk import AirflowFailException
+
+from unit.amazon.aws.utils.test_template_fields import validate_template_fields
+
+
+def _choose(conditions, if_ids="if_task", else_ids="else_task"):
+    """Instantiate with conditions list and call choose_branch."""
+    op = SageMakerConditionOperator(
+        task_id="test",
+        conditions=conditions,
+        if_task_ids=if_ids,
+        else_task_ids=else_ids,
+    )
+    return op.choose_branch(context=MagicMock(spec=dict))
+
+
+def _choose_flat(condition_type, left, right, if_ids="if_task", else_ids="else_task"):
+    """Instantiate with flat params and call choose_branch."""
+    op = SageMakerConditionOperator(
+        task_id="test",
+        condition_type=condition_type,
+        left_value=left,
+        right_value=right,
+        if_task_ids=if_ids,
+        else_task_ids=else_ids,
+    )
+    return op.choose_branch(context=MagicMock(spec=dict))
+
+
+def test_template_fields():
+    op = SageMakerConditionOperator(
+        task_id="test",
+        conditions=[{"type": "Equals", "left_value": 1, "right_value": 1}],
+        if_task_ids=["a"],
+        else_task_ids=["b"],
+    )
+    validate_template_fields(op)
+
+
+class TestConditionTypes:
+    """One true + one false per condition type, plus logical combinators."""
+
+    @pytest.mark.parametrize(
+        ("cond_type", "left", "right", "expected"),
+        [
+            ("Equals", 1, 1, "if"),
+            ("Equals", 1, 2, "else"),
+            ("GreaterThan", 5, 3, "if"),
+            ("GreaterThan", 3, 5, "else"),
+            ("GreaterThanOrEqualTo", 3, 3, "if"),
+            ("GreaterThanOrEqualTo", 2, 3, "else"),
+            ("LessThan", 3, 5, "if"),
+            ("LessThan", 5, 3, "else"),
+            ("LessThanOrEqualTo", 3, 3, "if"),
+            ("LessThanOrEqualTo", 5, 3, "else"),
+        ],
+    )
+    def test_comparison(self, cond_type, left, right, expected):
+        result = _choose([{"type": cond_type, "left_value": left, "right_value": right}])
+        assert result == (["if_task"] if expected == "if" else ["else_task"])
+
+    def test_in_true(self):
+        assert _choose([{"type": "In", "value": 1, "in_values": [1, 2]}]) == ["if_task"]
+
+    def test_in_false(self):
+        assert _choose([{"type": "In", "value": 4, "in_values": [1, 2]}]) == ["else_task"]
+
+    def test_not_negates(self):
+        cond = [{"type": "Not", "condition": {"type": "Equals", "left_value": 1, "right_value": 1}}]
+        assert _choose(cond) == ["else_task"]
+
+    def test_or_any_true(self):
+        cond = [
+            {
+                "type": "Or",
+                "conditions": [
+                    {"type": "Equals", "left_value": 1, "right_value": 2},
+                    {"type": "Equals", "left_value": 1, "right_value": 1},
+                ],
+            }
+        ]
+        assert _choose(cond) == ["if_task"]
+
+
+class TestAndSemantics:
+    def test_multiple_conditions_and(self):
+        """All true -> if, one false -> else."""
+        assert _choose(
+            [
+                {"type": "Equals", "left_value": 1, "right_value": 1},
+                {"type": "GreaterThan", "left_value": 5, "right_value": 3},
+            ]
+        ) == ["if_task"]
+        assert _choose(
+            [
+                {"type": "Equals", "left_value": 1, "right_value": 1},
+                {"type": "GreaterThan", "left_value": 2, "right_value": 10},
+            ]
+        ) == ["else_task"]
+
+
+class TestValueCasting:
+    def test_cast_numeric_and_passthrough(self):
+        """int string, float string, bool string, non-numeric string, non-string type."""
+        assert SageMakerConditionOperator._cast("42") == 42
+        assert SageMakerConditionOperator._cast("0.9") == 0.9
+        assert SageMakerConditionOperator._cast("true") is True
+        assert SageMakerConditionOperator._cast("us-east-1") == "us-east-1"
+        assert SageMakerConditionOperator._cast(42) == 42
+
+
+class TestValidation:
+    def test_empty_conditions_raises(self):
+        with pytest.raises(ValueError, match="At least 1 condition is required"):
+            SageMakerConditionOperator(task_id="t", conditions=[], if_task_ids="a", else_task_ids="b")
+
+    def test_unknown_type_raises(self):
+        with pytest.raises(ValueError, match="Unknown condition type"):
+            _choose([{"type": "FooBar", "left_value": 1, "right_value": 2}])
+
+    def test_type_mismatch_raises(self):
+        with pytest.raises(TypeError, match="Cannot compare"):
+            _choose([{"type": "GreaterThanOrEqualTo", "left_value": "hello", "right_value": 0.9}])
+
+    def test_none_operand_raises(self):
+        """Covers both Python None and Jinja-rendered string 'None'."""
+        with pytest.raises(TypeError, match="received None"):
+            _choose([{"type": "Equals", "left_value": None, "right_value": 1}])
+        with pytest.raises(TypeError, match="received None"):
+            _choose([{"type": "GreaterThanOrEqualTo", "left_value": "None", "right_value": 0.9}])
+
+    @pytest.mark.parametrize(
+        ("condition", "match_pattern"),
+        [
+            ({}, "missing required key 'type'"),
+            ({"type": "Equals", "left_value": 1}, "missing required key"),
+            ({"type": "Not"}, "missing required key"),
+            ({"type": "Or"}, "missing required key"),
+        ],
+    )
+    def test_missing_key_raises(self, condition, match_pattern):
+        with pytest.raises(ValueError, match=match_pattern):
+            _choose([condition])
+
+
+class TestFlatInterface:
+    def test_flat_condition(self):
+        """Flat params work for comparison and In types."""
+        assert _choose_flat("Equals", 1, 1) == ["if_task"]
+        assert _choose_flat("Equals", 1, 2) == ["else_task"]
+        assert _choose_flat("In", 1, [1, 2, 3]) == ["if_task"]
+        assert _choose_flat("In", 99, [1, 2, 3]) == ["else_task"]
+
+    def test_invalid_condition_type_raises(self):
+        with pytest.raises(ValueError, match="Unknown condition_type"):
+            SageMakerConditionOperator(
+                task_id="t",
+                condition_type="NotEquals",
+                left_value=1,
+                right_value=1,
+                if_task_ids="a",
+                else_task_ids="b",
+            )
+
+    def test_mutual_exclusion_raises(self):
+        with pytest.raises(ValueError, match="Cannot use 'condition_type' and 'conditions' together"):
+            SageMakerConditionOperator(
+                task_id="t",
+                condition_type="Equals",
+                left_value=1,
+                right_value=1,
+                conditions=[{"type": "Equals", "left_value": 1, "right_value": 1}],
+                if_task_ids="a",
+                else_task_ids="b",
+            )
+
+    def test_neither_provided_raises(self):
+        with pytest.raises(ValueError, match="Missing condition"):
+            SageMakerConditionOperator(task_id="t", if_task_ids="a", else_task_ids="b")
+
+    def test_optional_else_task_ids(self):
+        """else_task_ids defaults to empty list when omitted."""
+        op = SageMakerConditionOperator(
+            task_id="t",
+            conditions=[{"type": "Equals", "left_value": 1, "right_value": 1}],
+            if_task_ids=["deploy"],
+        )
+        assert op.else_task_ids == []
+
+    def test_no_else_branch_raises_on_false(self):
+        """When else_task_ids is empty and conditions are false, raises AirflowFailException."""
+        op = SageMakerConditionOperator(
+            task_id="check",
+            conditions=[{"type": "Equals", "left_value": 1, "right_value": 2}],
+            if_task_ids=["deploy"],
+        )
+        with pytest.raises(AirflowFailException, match="Condition check failed in task 'check'"):
+            op.choose_branch(context=MagicMock(spec=dict))
