[flyteagent] FileSensor Timeout (Remote Execution Only) Enhancement (#3191)

Signed-off-by: Future-Outlier <[email protected]>

Author: Future-Outlier

Diff for: flytekit/sensor/base_sensor.py

@@ -63,10 +63,15 @@ def __init__(
             annotation = type_hints.get(k, None)
             inputs[k] = annotation
 
-        if kwargs.get("metadata", None) and timeout:
-            raise ValueError("You cannot set timeout and metadata at the same time in the sensor")
-
-        metadata = TaskMetadata(timeout=timeout)
+        # Handle metadata and timeout logic
+        metadata = kwargs.pop("metadata", None)
+        if metadata is not None and timeout is not None:
+            if metadata.timeout is not None:
+                raise ValueError("You cannot set both timeout and metadata parameters at the same time in the sensor")
+            else:
+                metadata.timeout = timeout
+        else:
+            metadata = TaskMetadata(timeout=timeout) if timeout else TaskMetadata()
 
         super().__init__(
             task_type=task_type,

Diff for: tests/flytekit/unit/sensor/test_file_sensor.py

@@ -79,8 +79,19 @@ def test_agent_executor_timeout_logging():
     assert sensor.metadata.timeout == datetime.timedelta(seconds=60)
 
 def test_file_sensor_set_timeout_and_metadata_at_the_same_time():
+    """
+    The test is the combination of the following cases:
+    1. Set timeout and metadata with timeout at the same time
+    2. Set timeout and metadata without timeout
+    3. Set timeout and no metadata
+    4. Not set timeout and metadata with timeout
+    """
+
     timeout = datetime.timedelta(seconds=60)
-    metadata = TaskMetadata(timeout=timeout)
 
-    with pytest.raises(ValueError, match="You cannot set timeout and metadata at the same time in the sensor"):
-        FileSensor(name="test_sensor", timeout=60, metadata=metadata)
+    with pytest.raises(ValueError, match="You cannot set both timeout and metadata parameters at the same time in the sensor"):
+        FileSensor(name="test_sensor", timeout=60, metadata=TaskMetadata(timeout=timeout))
+
+    FileSensor(name="test_sensor", timeout=60, metadata=TaskMetadata())
+    FileSensor(name="test_sensor", timeout=60)
+    FileSensor(name="test_sensor", metadata=TaskMetadata(timeout=timeout))