amoghrajesh commented on code in PR #54212:
URL: https://github.com/apache/airflow/pull/54212#discussion_r2262114522


##########
airflow-core/tests/unit/models/test_dagbag.py:
##########
@@ -77,6 +77,17 @@ def setup_class(self):
     def teardown_class(self):
         db_clean_up()
 
+    def test_timeout_context_manager_raises_exception(self):
+        """Test that the timeout context manager raises AirflowTaskTimeout 
when time limit is exceeded."""
+        import time
+
+        from airflow.exceptions import AirflowTaskTimeout
+        from airflow.models.dagbag import timeout
+
+        with pytest.raises(AirflowTaskTimeout):
+            with timeout(1, "Test timeout"):
+                time.sleep(2)

Review Comment:
   Added a simple test for basic functionality too.



-- 
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]

Reply via email to