ferruzzi commented on code in PR #57222:
URL: https://github.com/apache/airflow/pull/57222#discussion_r2466839837
##########
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):
Review Comment:
Yeah, that's fair. I wanted to make sure that the user-experience is the
same for the custom vs built-in because in my initial implementation the custom
references were `ReferenceModels.MyCustomRef` and the built-ins were at
`DeadlineReference.SomeBuiltinRef` so I added this, but if you think it's
overkill, I can definitely drop it.
--
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]
