o-nikolas commented on code in PR #57222:
URL: https://github.com/apache/airflow/pull/57222#discussion_r2467207062
##########
airflow-core/tests/unit/models/test_deadline.py:
##########
@@ -583,3 +593,272 @@ def test_deadline_reference_creation(self):
custom_reference = DeadlineReference.AVERAGE_RUNTIME(max_runs=5,
min_runs=3)
assert custom_reference.max_runs == 5
assert custom_reference.min_runs == 3
+
+
+class TestCustomDeadlineReference:
+ class MyCustomRef(ReferenceModels.BaseDeadlineReference):
+ def _evaluate_with(self, *, session: Session, **kwargs) -> datetime:
+ return timezone.datetime(DEFAULT_DATE)
+
+ class MyInvalidCustomRef:
+ pass
+
+ class MyCustomRefWithKwargs(ReferenceModels.BaseDeadlineReference):
+ required_kwargs = {"custom_id"}
+
+ def _evaluate_with(self, *, session: Session, **kwargs) -> datetime:
+ return timezone.datetime(DEFAULT_DATE)
+
+ def setup_method(self):
+ """Store original state before each test."""
+ self.original_dagrun_created = DeadlineReference.TYPES.DAGRUN_CREATED
+ self.original_dagrun_queued = DeadlineReference.TYPES.DAGRUN_QUEUED
+ self.original_dagrun = DeadlineReference.TYPES.DAGRUN
+ self.original_attrs = set(dir(ReferenceModels))
+ self.original_deadline_attrs = set(dir(DeadlineReference))
+
+ def teardown_method(self):
+ """Restore original TYPES and attrs after each test."""
+ DeadlineReference.TYPES.DAGRUN_CREATED = self.original_dagrun_created
+ DeadlineReference.TYPES.DAGRUN_QUEUED = self.original_dagrun_queued
+ DeadlineReference.TYPES.DAGRUN = self.original_dagrun
+
+ # Clean up ReferenceModels attributes.
+ for attr in set(dir(ReferenceModels)):
+ if attr not in self.original_attrs:
+ delattr(ReferenceModels, attr)
+
+ # Clean up DeadlineReference attributes.
+ for attr in set(dir(DeadlineReference)):
+ if attr not in self.original_deadline_attrs:
+ delattr(DeadlineReference, attr)
+
+ def test_custom_reference_consistent_access_pattern(self):
+ DeadlineReference.register_custom_reference(self.MyCustomRef)
+
+ # Should be accessible through DeadlineReference like built-ins.
+ custom_ref = getattr(DeadlineReference, self.MyCustomRef.__name__)
+ assert custom_ref.__class__ is self.MyCustomRef
+
+ # Should behave like built-in references.
+ assert hasattr(custom_ref, "_evaluate_with")
+ assert callable(custom_ref._evaluate_with)
+
+ def test_register_custom_reference_default_timing(self):
+ result = DeadlineReference.register_custom_reference(self.MyCustomRef)
+
+ # Should return the class.
+ assert result is self.MyCustomRef
+
+ # Should be registered with ReferenceModels.
+ assert hasattr(ReferenceModels, self.MyCustomRef.__name__)
+ assert getattr(ReferenceModels, self.MyCustomRef.__name__) is
self.MyCustomRef
+
+ # Should be accessible through DeadlineReference.
+ assert hasattr(DeadlineReference, self.MyCustomRef.__name__)
+ assert getattr(DeadlineReference, self.MyCustomRef.__name__).__class__
is self.MyCustomRef
+ assert self.MyCustomRef in DeadlineReference.TYPES.DAGRUN_CREATED
+ assert self.MyCustomRef not in DeadlineReference.TYPES.DAGRUN_QUEUED
+
+ # Should update combined DAGRUN tuple
+ assert self.MyCustomRef in DeadlineReference.TYPES.DAGRUN
+
+ def
test_register_custom_reference_dagrun_created_with_explicit_timing(self):
+ result = DeadlineReference.register_custom_reference(
+ self.MyCustomRef, DeadlineReference.TYPES.DAGRUN_CREATED
+ )
+
+ assert result is self.MyCustomRef
+ assert hasattr(ReferenceModels, self.MyCustomRef.__name__)
+ assert self.MyCustomRef in DeadlineReference.TYPES.DAGRUN_CREATED
+ assert self.MyCustomRef not in DeadlineReference.TYPES.DAGRUN_QUEUED
+ assert self.MyCustomRef in DeadlineReference.TYPES.DAGRUN
+
+ def test_register_custom_reference_with_dagrun_queued(self):
+ result = DeadlineReference.register_custom_reference(
+ self.MyCustomRef, DeadlineReference.TYPES.DAGRUN_QUEUED
+ )
+
+ assert result is self.MyCustomRef
+ assert hasattr(ReferenceModels, self.MyCustomRef.__name__)
+ assert self.MyCustomRef not in DeadlineReference.TYPES.DAGRUN_CREATED
+ assert self.MyCustomRef in DeadlineReference.TYPES.DAGRUN_QUEUED
+ assert self.MyCustomRef in DeadlineReference.TYPES.DAGRUN
+
+ def test_register_custom_reference_invalid_inheritance(self):
+ with pytest.raises(ValueError, match="InvalidCustomRef must inherit
from BaseDeadlineReference"):
+
DeadlineReference.register_custom_reference(self.MyInvalidCustomRef)
+
+ def test_register_custom_reference_invalid_timing(self):
+ invalid_timing = ("not", "a", "valid", "timing")
+
+ with pytest.raises(
+ ValueError, match="Invalid timing value; must be a valid
DeadlineReference.TYPES option"
+ ):
+ DeadlineReference.register_custom_reference(self.MyCustomRef,
invalid_timing)
+
+ def test_register_custom_reference_with_required_kwargs(self):
+ result =
DeadlineReference.register_custom_reference(self.MyCustomRefWithKwargs)
+
+ assert result is self.MyCustomRefWithKwargs
+ assert hasattr(ReferenceModels, self.MyCustomRefWithKwargs.__name__)
+ assert self.MyCustomRefWithKwargs in
DeadlineReference.TYPES.DAGRUN_CREATED
+
+ def test_register_multiple_custom_references(self):
+ class TestCustomRef1(ReferenceModels.BaseDeadlineReference):
+ def _evaluate_with(self, *, session: Session, **kwargs) ->
datetime:
+ return timezone.datetime(DEFAULT_DATE)
+
+ class TestCustomRef2(ReferenceModels.BaseDeadlineReference):
+ def _evaluate_with(self, *, session: Session, **kwargs) ->
datetime:
+ return timezone.datetime(DEFAULT_DATE)
+
+ # Register first reference to DAGRUN_CREATED
Review Comment:
Up to you, happy to disagree and commit 👍
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
