langfun 0.1.2.dev202510310805__py3-none-any.whl → 0.1.2.dev202511020804__py3-none-any.whl

langfun/env/base_feature_test.py

@@ -0,0 +1,228 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed 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.
+import unittest
+
+from langfun.env import test_utils
+
+TestingEnvironment = test_utils.TestingEnvironment
+TestingNonSandboxBasedFeature = test_utils.TestingNonSandboxBasedFeature
+TestingEventHandler = test_utils.TestingEventHandler
+
+
+class NonSandboxBasedFeatureTests(unittest.TestCase):
+
+  def test_basics(self):
+    feature = TestingNonSandboxBasedFeature()
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={'test_feature': feature},
+        event_handler=event_handler,
+    )
+    self.assertFalse(env.is_online)
+    self.assertEqual(len(list(env.non_sandbox_based_features())), 1)
+    with env:
+      self.assertTrue(env.is_online)
+      with env.test_feature('session1') as feature:
+        self.assertIsNone(feature.sandbox)
+        self.assertEqual(feature.session_id, 'session1')
+
+    self.assertEqual(
+        event_handler.logs,
+        [
+            '[testing-env/test_feature] feature setup',
+            '[testing-env] environment started',
+            '[testing-env/test_feature@session1] feature setup session',
+            '[testing-env/test_feature@session1] feature teardown session',
+            '[testing-env/test_feature] feature teardown',
+            '[testing-env] environment shutdown'
+        ]
+    )
+
+  def test_feature_setup_error(self):
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_setup_error=ValueError
+            )
+        },
+        event_handler=event_handler,
+    )
+    with self.assertRaises(ValueError):
+      with env:
+        pass
+    self.assertEqual(
+        event_handler.logs,
+        [
+            '[testing-env/test_feature] feature setup with ValueError',
+            '[testing-env] environment started with ValueError',
+            '[testing-env/test_feature] feature teardown',
+            '[testing-env] environment shutdown'
+        ]
+    )
+
+  def test_feature_teardown_error(self):
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_teardown_error=ValueError
+            )
+        },
+        event_handler=event_handler,
+    )
+    with env:
+      pass
+    self.assertEqual(
+        event_handler.logs,
+        [
+            '[testing-env/test_feature] feature setup',
+            '[testing-env] environment started',
+            '[testing-env/test_feature] feature teardown with ValueError',
+            '[testing-env] environment shutdown'
+        ]
+    )
+
+  def test_feature_setup_session_error(self):
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_setup_session_error=ValueError
+            )
+        },
+        event_handler=event_handler,
+    )
+    with env:
+      with self.assertRaises(ValueError):
+        with env.test_feature('session1'):
+          pass
+    self.assertEqual(
+        event_handler.logs,
+        [
+            # pylint: disable=line-too-long
+            '[testing-env/test_feature] feature setup',
+            '[testing-env] environment started',
+            '[testing-env/test_feature@session1] feature setup session with ValueError',
+            '[testing-env/test_feature@session1] feature teardown session',
+            '[testing-env/test_feature] feature teardown',
+            '[testing-env] environment shutdown',
+            # pylint: enable=line-too-long
+        ]
+    )
+
+  def test_feature_teardown_session_error(self):
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_teardown_session_error=ValueError
+            )
+        },
+        event_handler=event_handler,
+    )
+    with env:
+      with env.test_feature('session1'):
+        pass
+    self.assertEqual(
+        event_handler.logs,
+        [
+            # pylint: disable=line-too-long
+            '[testing-env/test_feature] feature setup',
+            '[testing-env] environment started',
+            '[testing-env/test_feature@session1] feature setup session',
+            '[testing-env/test_feature@session1] feature teardown session with ValueError',
+            '[testing-env/test_feature] feature teardown',
+            '[testing-env] environment shutdown',
+            # pylint: enable=line-too-long
+        ]
+    )
+
+  def test_feature_housekeeping(self):
+    event_handler = TestingEventHandler(
+        log_sandbox_status=False,
+        log_feature_setup=False,
+        log_housekeep=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                housekeep_interval=0.1
+            )
+        },
+        event_handler=event_handler,
+        housekeep_interval=0.2
+    )
+    with env:
+      env.wait_for_housekeeping()
+    self.assertIn(
+        '[testing-env/test_feature] feature housekeeping 0',
+        event_handler.logs
+    )
+
+  def test_feature_housekeeping_error(self):
+    event_handler = TestingEventHandler(
+        log_sandbox_status=False,
+        log_feature_setup=False,
+        log_housekeep=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_housekeep_error=ValueError,
+                housekeep_interval=0.1
+            )
+        },
+        event_handler=event_handler,
+        housekeep_interval=0.2
+    )
+    with env:
+      env.wait_for_housekeeping()
+      self.assertFalse(env.is_online)
+    self.assertIn(
+        '[testing-env/test_feature] feature housekeeping 0 with ValueError',
+        event_handler.logs
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()
+

langfun/env/base_sandbox.py

@@ -29,7 +29,6 @@ import time
 from typing import Annotated, Any, Iterator
 
 from langfun.env import interface
-from langfun.env.event_handlers import base as event_handler_base
 import pyglove as pg
 
 
@@ -198,7 +197,7 @@ class BaseSandbox(interface.Sandbox):
         for name, feature in self.environment.features.items()
         if feature.is_applicable(self.image_id)
     })
-    self._event_handlers = []
+    self._event_handler = self.environment.event_handler
     self._enable_pre_session_setup = (
         self.reusable and self.proactive_session_setup
     )
@@ -247,20 +246,6 @@ class BaseSandbox(interface.Sandbox):
     """Returns the housekeeping counter."""
     return self._housekeep_counter
 
-  def add_event_handler(
-      self,
-      event_handler: event_handler_base.EventHandler | None
-  ) -> None:
-    """Sets the event handler for the sandbox."""
-    self._event_handlers.append(event_handler)
-
-  def remove_event_handler(
-      self,
-      event_handler: event_handler_base.EventHandler | None
-  ) -> None:
-    """Removes the event handler for the sandbox."""
-    self._event_handlers.remove(event_handler)
-
   @property
   def state_errors(self) -> list[interface.SandboxStateError]:
     """Returns all errors encountered during sandbox lifecycle."""
@@ -648,7 +633,6 @@ class BaseSandbox(interface.Sandbox):
       shutdown_sandbox = True
 
     self._session_start_time = None
-    self._session_event_handler = None
 
     if shutdown_sandbox:
       self.shutdown()
@@ -664,7 +648,6 @@ class BaseSandbox(interface.Sandbox):
   def track_activity(
       self,
       name: str,
-      feature: interface.Feature | None = None,
       **kwargs: Any
   ) -> Iterator[None]:
     """Tracks an activity for the sandbox."""
@@ -678,7 +661,6 @@ class BaseSandbox(interface.Sandbox):
     finally:
       self.on_activity(
           name=name,
-          feature=feature,
           duration=time.time() - start_time,
           error=error,
           **kwargs
@@ -771,8 +753,7 @@ class BaseSandbox(interface.Sandbox):
       error: BaseException | None = None
   ) -> None:
     """Called when the sandbox is started."""
-    for handler in self._event_handlers:
-      handler.on_sandbox_start(self.environment, self, duration, error)
+    self._event_handler.on_sandbox_start(self, duration, error)
 
   def on_status_change(
       self,
@@ -781,14 +762,9 @@ class BaseSandbox(interface.Sandbox):
   ) -> None:
     """Called when the sandbox status changes."""
     status_duration = time.time() - self._status_start_time
-    for handler in self._event_handlers:
-      handler.on_sandbox_status_change(
-          self.environment,
-          self,
-          old_status,
-          new_status,
-          status_duration
-      )
+    self._event_handler.on_sandbox_status_change(
+        self, old_status, new_status, status_duration
+    )
 
   def on_shutdown(
       self,
@@ -796,14 +772,13 @@ class BaseSandbox(interface.Sandbox):
       error: BaseException | None = None
   ) -> None:
     """Called when the sandbox is shutdown."""
-    if self._start_time is None:
-      lifetime = 0.0
-    else:
-      lifetime = time.time() - self._start_time
-    for handler in self._event_handlers:
-      handler.on_sandbox_shutdown(
-          self.environment, self, duration, lifetime, error
-      )
+    self._event_handler.on_sandbox_shutdown(
+        sandbox=self,
+        duration=duration,
+        lifetime=(0.0 if self._start_time is None
+                  else (time.time() - self._start_time)),
+        error=error
+    )
 
   def on_housekeep(
       self,
@@ -812,72 +787,13 @@ class BaseSandbox(interface.Sandbox):
       **kwargs
   ) -> None:
     """Called when the sandbox finishes a round of housekeeping."""
-    counter = self._housekeep_counter
-    for handler in self._event_handlers:
-      handler.on_sandbox_housekeep(
-          self.environment, self, counter, duration, error, **kwargs
-      )
-
-  def on_feature_setup(
-      self,
-      feature: interface.Feature,
-      duration: float,
-      error: BaseException | None = None
-  ) -> None:
-    """Called when a feature is setup."""
-    for handler in self._event_handlers:
-      handler.on_feature_setup(
-          self.environment, self, feature, duration, error
-      )
-
-  def on_feature_teardown(
-      self,
-      feature: interface.Feature,
-      duration: float,
-      error: BaseException | None = None
-  ) -> None:
-    """Called when a feature is teardown."""
-    for handler in self._event_handlers:
-      handler.on_feature_teardown(
-          self.environment, self, feature, duration, error
-      )
-
-  def on_feature_setup_session(
-      self,
-      feature: interface.Feature,
-      duration: float,
-      error: BaseException | None = None
-  ) -> None:
-    """Called when a feature is setup for a user session."""
-    for handler in self._event_handlers:
-      handler.on_feature_setup_session(
-          self.environment, self, feature, self.session_id, duration, error
-      )
-
-  def on_feature_teardown_session(
-      self,
-      feature: interface.Feature,
-      duration: float,
-      error: BaseException | None = None
-  ) -> None:
-    """Called when a feature is teardown for a user session."""
-    for handler in self._event_handlers:
-      handler.on_feature_teardown_session(
-          self.environment, self, feature, self.session_id, duration, error
-      )
-
-  def on_feature_housekeep(
-      self,
-      feature: interface.Feature,
-      counter: int,
-      duration: float,
-      error: BaseException | None = None
-  ) -> None:
-    """Called when a feature is housekeeping."""
-    for handler in self._event_handlers:
-      handler.on_feature_housekeep(
-          self.environment, self, feature, counter, duration, error
-      )
+    self._event_handler.on_sandbox_housekeep(
+        sandbox=self,
+        counter=self._housekeep_counter,
+        duration=duration,
+        error=error,
+        **kwargs
+    )
 
   def on_session_start(
       self,
@@ -886,31 +802,29 @@ class BaseSandbox(interface.Sandbox):
       error: BaseException | None = None
   ) -> None:
     """Called when the user session starts."""
-    for handler in self._event_handlers:
-      handler.on_session_start(
-          self.environment, self, session_id, duration, error
-      )
+    self._event_handler.on_sandbox_session_start(
+        sandbox=self,
+        session_id=session_id,
+        duration=duration,
+        error=error
+    )
 
   def on_activity(
       self,
       name: str,
       duration: float,
-      feature: interface.Feature | None = None,
       error: BaseException | None = None,
       **kwargs
   ) -> None:
     """Called when a sandbox activity is performed."""
-    for handler in self._event_handlers:
-      handler.on_sandbox_activity(
-          name=name,
-          environment=self.environment,
-          sandbox=self,
-          feature=feature,
-          session_id=self.session_id,
-          duration=duration,
-          error=error,
-          **kwargs
-      )
+    self._event_handler.on_sandbox_activity(
+        name=name,
+        sandbox=self,
+        session_id=self.session_id,
+        duration=duration,
+        error=error,
+        **kwargs
+    )
 
   def on_session_end(
       self,
@@ -919,8 +833,10 @@ class BaseSandbox(interface.Sandbox):
       error: BaseException | None = None
   ) -> None:
     """Called when the user session ends."""
-    lifetime = time.time() - self._session_start_time
-    for handler in self._event_handlers:
-      handler.on_session_end(
-          self.environment, self, session_id, duration, lifetime, error
-      )
+    self._event_handler.on_sandbox_session_end(
+        sandbox=self,
+        session_id=session_id,
+        duration=duration,
+        lifetime=time.time() - self._session_start_time,
+        error=error
+    )