langfun 0.1.2.dev202509240805__py3-none-any.whl → 0.1.2.dev202509260805__py3-none-any.whl

langfun/env/test_utils.py

@@ -0,0 +1,459 @@
+# 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.
+"""Test utils for base environment."""
+
+import contextlib
+import time
+from typing import Iterator, Type
+
+from langfun.env import base_environment
+from langfun.env import base_feature
+from langfun.env import base_sandbox
+from langfun.env import interface
+from langfun.env.event_handlers import base as event_handler_base
+import pyglove as pg
+
+
+class TestingEnvironment(base_environment.BaseEnvironment):
+  """Testing environment for unit tests."""
+
+  housekeep_interval: float = 0.0
+  simulate_start_error: Type[BaseException] | None = None
+  simulate_shutdown_error: Type[BaseException] | None = None
+  simulate_ping_error: Type[BaseException] | None = None
+  offline: bool = False
+
+  @property
+  def id(self) -> interface.Environment.Id:
+    return interface.Environment.Id('testing-env')
+
+  def wait_for_housekeeping(self):
+    housekeep_counter = self.housekeep_counter
+    while self.housekeep_counter == housekeep_counter:
+      time.sleep(0.01)
+
+  def _create_sandbox(
+      self,
+      sandbox_id: str,
+      reusable: bool,
+      proactive_session_setup: bool,
+      keepalive_interval: float | None,
+  ) -> base_sandbox.BaseSandbox:
+    return TestingSandbox(
+        environment=self,
+        id=interface.Sandbox.Id(
+            environment_id=self.id,
+            sandbox_id=sandbox_id
+        ),
+        reusable=reusable,
+        proactive_session_setup=proactive_session_setup,
+        keepalive_interval=keepalive_interval,
+        simulate_start_error=self.simulate_start_error,
+        simulate_shutdown_error=self.simulate_shutdown_error,
+        simulate_ping_error=self.simulate_ping_error,
+    )
+
+
+class TestingSandbox(base_sandbox.BaseSandbox):
+  """Testing sandbox for unit tests."""
+
+  simulate_start_error: Type[BaseException] | None = None
+  simulate_shutdown_error: Type[BaseException] | None = None
+  simulate_ping_error: Type[BaseException] | None = None
+
+  def _on_bound(self) -> None:
+    super()._on_bound()
+    self._shell_history = []
+    self._ping_history = []
+
+  def _raise_error(self, message, error_type: Type[BaseException], **kwargs):
+    if (error_type is interface.SandboxStateError or
+        issubclass(error_type, interface.SandboxStateError)):
+      kwargs['sandbox'] = self
+      raise error_type(message, **kwargs)
+    else:
+      raise error_type(message)
+
+  def wait_until_not(
+      self,
+      status: interface.Sandbox.Status | tuple[interface.Sandbox.Status, ...]
+  ) -> None:
+    if not isinstance(status, tuple):
+      status = (status,)
+    while self.status in status:
+      time.sleep(0.01)
+
+  def wait_until_next_housekeep(self) -> None:
+    housekeep_counter = self.housekeep_counter
+    while self.housekeep_counter == housekeep_counter:
+      time.sleep(0.01)
+
+  def _start(self) -> None:
+    if self.simulate_start_error:
+      self._raise_error('Sandbox start error', self.simulate_start_error)
+    super()._start()
+
+  def _shutdown(self) -> None:
+    if self.simulate_shutdown_error:
+      self._raise_error('Sandbox shutdown error', self.simulate_shutdown_error)
+    super()._shutdown()
+
+  @base_sandbox.sandbox_service(critical_errors=(RuntimeError,))
+  def shell(
+      self,
+      code: str,
+      raise_error: Type[BaseException] | None = None,
+  ) -> str:
+    self._shell_history.append(code)
+    if raise_error is not None:
+      self._raise_error(f'shell "{code}" failed', raise_error)
+    return f'shell "{code}" succeeded'
+
+  def _ping(self) -> None:
+    self._ping_history.append(not self.simulate_ping_error)
+    if self.simulate_ping_error:
+      self._raise_error('Ping error', self.simulate_ping_error, code='ping')
+
+
+class TestingFeature(base_feature.BaseFeature):
+  """Testing feature for unit tests."""
+
+  housekeep_interval = 0
+  setup_session_delay: float = 0.0
+  simulate_housekeep_error: Type[BaseException] | None = None
+  simulate_setup_error: Type[BaseException] | None = None
+  simulate_teardown_error: Type[BaseException] | None = None
+  simulate_setup_session_error: Type[BaseException] | None = None
+  simulate_teardown_session_error: Type[BaseException] | None = None
+  call_end_session_on_teardown_session: bool = False
+
+  class Service:
+    """Sandbox."""
+
+    def __init__(self, sandbox: interface.Sandbox):
+      self._sandbox = sandbox
+
+    def do(self, code: str, raise_error: Type[BaseException] | None = None):
+      self._sandbox.shell(code, raise_error=raise_error)
+
+  def _raise_error(self, message, error_type: Type[BaseException], **kwargs):
+    self._sandbox._raise_error(message, error_type, **kwargs)  # pylint: disable=protected-access
+
+  def _setup(self) -> None:
+    if self.simulate_setup_error:
+      self._raise_error(f'{self.name} setup error', self.simulate_setup_error)
+    self.sandbox.shell(f'"{self.name}" setup')
+
+  def _teardown(self) -> None:
+    if self.simulate_teardown_error:
+      self._raise_error(
+          f'{self.name} teardown error', self.simulate_teardown_error
+      )
+    self.sandbox.shell(f'"{self.name}" teardown')
+
+  def _setup_session(self) -> None:
+    if self.setup_session_delay > 0:
+      time.sleep(self.setup_session_delay)
+
+    if self.simulate_setup_session_error:
+      self._raise_error(
+          'Feature session setup error', self.simulate_setup_session_error
+      )
+    self.sandbox.shell(f'"{self.name}" setup session')
+
+  def _teardown_session(self) -> None:
+    if self.simulate_teardown_session_error:
+      self._raise_error(
+          'Feature session teardown error', self.simulate_teardown_session_error
+      )
+    self.sandbox.shell(f'"{self.name}" teardown session')
+    if self.call_end_session_on_teardown_session:
+      self.sandbox.end_session()
+
+  @base_sandbox.sandbox_service()
+  def num_shell_calls(self) -> int:
+    return len(self.sandbox._shell_history)  # pylint: disable=protected-access
+
+  @base_sandbox.sandbox_service()
+  def bad_shell_call(self) -> None:
+    self.sandbox.shell('bad command', raise_error=RuntimeError)
+
+  @base_sandbox.sandbox_service()
+  def show_session_id(self):
+    return self.session_id
+
+  @base_sandbox.sandbox_service()
+  def call_with_varargs(self, code: str, *args, **kwargs):
+    del code, args, kwargs
+    return 0
+
+  def _on_bound(self) -> None:
+    super()._on_bound()
+    self._service = None
+
+  @base_sandbox.sandbox_service()
+  @contextlib.contextmanager
+  def my_service(self) -> Iterator[Service]:
+    try:
+      self._service = TestingFeature.Service(sandbox=self.sandbox)
+      yield self._service
+    finally:
+      self._service = None
+
+  def _housekeep(self) -> None:
+    if self.simulate_housekeep_error:
+      raise interface.SandboxStateError(
+          'House keeping error', sandbox=self.sandbox
+      )
+
+
+class TestingEventHandler(
+    pg.Object, event_handler_base.EventHandler
+):
+  """Testing environment event handler for unit tests."""
+
+  log_sandbox_status: bool = False
+  log_feature_setup: bool = True
+  log_session_setup: bool = False
+  log_housekeep: bool = False
+
+  def _on_bound(self) -> None:
+    super()._on_bound()
+    self._logs = []
+
+  @property
+  def logs(self) -> list[str]:
+    return self._logs
+
+  def _add_message(self, message: str, error: BaseException | None) -> None:
+    """Adds a message to the history."""
+    if error is None:
+      self._logs.append(message)
+    else:
+      self._logs.append(f'{message} with {error.__class__.__name__}')
+
+  def on_environment_start(
+      self,
+      environment: interface.Environment,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when the environment is started."""
+    assert duration > 0
+    self._add_message(f'[{environment.id}] environment started', error)
+
+  def on_environment_housekeep(
+      self,
+      environment: interface.Environment,
+      counter: int,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when the environment finishes a round of housekeeping."""
+    assert duration > 0
+    if self.log_housekeep:
+      self._add_message(
+          f'[{environment.id}] environment housekeeping {counter}', error
+      )
+
+  def on_environment_shutdown(
+      self,
+      environment: interface.Environment,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when the environment is shutdown."""
+    assert lifetime is not None
+    self._add_message(f'[{environment.id}] environment shutdown', error)
+
+  def on_sandbox_start(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    assert duration > 0
+    self._add_message(f'[{sandbox.id}] sandbox started', error)
+
+  def on_sandbox_status_change(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      old_status: interface.Sandbox.Status,
+      new_status: interface.Sandbox.Status,
+      span: float
+  ) -> None:
+    assert span > 0
+    if self.log_sandbox_status:
+      self._add_message(
+          f'[{sandbox.id}] {old_status.value} -> {new_status.value}',
+          None,
+      )
+
+  def on_sandbox_shutdown(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    assert lifetime is not None
+    self._add_message(f'[{sandbox.id}] sandbox shutdown', error)
+
+  def on_sandbox_housekeep(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      counter: int,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    assert duration > 0
+    if self.log_housekeep:
+      self._add_message(
+          f'[{sandbox.id}] sandbox housekeeping {counter}', error
+      )
+
+  def on_feature_setup(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      feature: interface.Feature,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is setup."""
+    assert duration > 0
+    if self.log_feature_setup:
+      self._add_message(
+          f'[{sandbox.id}/{feature.name}] feature setup', error
+      )
+
+  def on_feature_teardown(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      feature: interface.Feature,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is teardown."""
+    assert duration > 0
+    if self.log_feature_setup:
+      self._add_message(
+          f'[{sandbox.id}/{feature.name}] feature teardown', error
+      )
+
+  def on_feature_setup_session(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      feature: interface.Feature,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is setup."""
+    assert duration > 0
+    if self.log_session_setup:
+      self._add_message(
+          f'[{sandbox.id}/{feature.name}] feature setup session', error
+      )
+
+  def on_feature_teardown_session(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      feature: interface.Feature,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is teardown."""
+    assert duration > 0
+    if self.log_session_setup:
+      self._add_message(
+          f'[{sandbox.id}/{feature.name}] feature teardown session', error
+      )
+
+  def on_feature_housekeep(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      feature: interface.Feature,
+      counter: int,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is housekeeping."""
+    assert duration > 0
+    if self.log_housekeep:
+      self._add_message(
+          f'[{sandbox.id}/{feature.name}] feature housekeeping {counter}', error
+      )
+
+  def on_session_start(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session starts."""
+    assert duration > 0
+    self._add_message(
+        f'[{sandbox.id}] session {session_id!r} started', error
+    )
+
+  def on_session_end(
+      self,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session ends."""
+    assert lifetime > 0
+    self._add_message(
+        f'[{sandbox.id}] session {session_id!r} ended', error
+    )
+
+  def on_sandbox_activity(
+      self,
+      name: str,
+      environment: interface.Environment,
+      sandbox: interface.Sandbox,
+      feature: interface.Feature | None,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None,
+      *,
+      code: str | None = None,
+      **kwargs
+  ) -> None:
+    """Called when a sandbox activity is performed."""
+    del environment, kwargs
+    if session_id is None:
+      log_id = sandbox.id
+    else:
+      log_id = f'{sandbox.id}/{session_id}'
+
+    if feature is not None:
+      log_id = f'{log_id}/{feature.name}'
+    self._add_message(
+        f'[{log_id}] {name}: {code}', error
+    )

{langfun-0.1.2.dev202509240805.dist-info → langfun-0.1.2.dev202509260805.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfun
-Version: 0.1.2.dev202509240805
+Version: 0.1.2.dev202509260805
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202509240805.dist-info → langfun-0.1.2.dev202509260805.dist-info}/RECORD

@@ -165,17 +165,22 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun/env/__init__.py,sha256=WOht-DaYtVmO8U-viQbWmCUCCPcJwBELBiQYUfGAG5w,1532
-langfun/env/base_environment.py,sha256=73d6hs2XhJd2IK--H7Pfs-1PG5TfJ89x_yk94doq0G0,16465
-langfun/env/base_feature.py,sha256=lphS2GmwCYR-ypPGIzn0KcLzxUcO2l9R2ZbJqcqOkOI,5794
-langfun/env/base_sandbox.py,sha256=c-HskeubKOH-P2ZrQG3HFKkhtX3ZADYj4QIMemQNGbo,32202
-langfun/env/base_test.py,sha256=CIzZjfSkZEDsHARF4aoOZfq5T8DN8i8AW-swFsFw0Gk,60328
-langfun/env/interface.py,sha256=wOu_yQO8uEt5lU88FeqqaInBLg7WgYYtVeGtPsanE90,29031
-langfun/env/interface_test.py,sha256=nrqFLtpHd0vnOjoIX9cjV5WD1c-7J3WW29aogNu_94A,1395
+langfun/env/__init__.py,sha256=m6Y8y16ms9lytvO_r-Br8HmTp2rNDhb3R6JJaH5dEEk,1491
+langfun/env/base_environment.py,sha256=Vjpmck3k3UE3KaHq69dIuGbevbRJFlY6PvQEeZ18MtY,18087
+langfun/env/base_feature.py,sha256=dztFq1lcdXgUh3m8pNRC3vruyIdHCob27reiXz_02q8,6436
+langfun/env/base_sandbox.py,sha256=tR__USJZojDrGtnhmDqkbCHrvcrNUDrTwc9mpStyYZ0,36721
+langfun/env/base_test.py,sha256=cBZEcoJG16L_XLVyTgSCxUxy1cukrkpSNNEgiM5LZb4,59137
+langfun/env/interface.py,sha256=2Amf-_op7dGRF8c4-wYxcFxs1UCBYz1AB20Lk7__V4E,25724
+langfun/env/interface_test.py,sha256=hfQn4RRTEo1YfVHXTPzH1puzD14BTo8R_5v1IpXVZ90,1398
 langfun/env/load_balancers.py,sha256=qRhCthqzjZIQBwta8qC1C0s0J-VQAVomJQqI7Nqv-r4,1948
-langfun/env/load_balancers_test.py,sha256=gJKJ3K5_WS7hsz9Jv3nPEjlb04Z6MmrZpfCu86siYKU,3967
-langfun-0.1.2.dev202509240805.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202509240805.dist-info/METADATA,sha256=r1KTHcsTEVX6kFjs88_yhZ3xY-k0urqHENZadNUnyi8,7380
-langfun-0.1.2.dev202509240805.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-langfun-0.1.2.dev202509240805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202509240805.dist-info/RECORD,,
+langfun/env/load_balancers_test.py,sha256=KIDIwZeQHcW0fCL5ScZD0lm0QV4_X2y-rMqdd6R1gLc,3737
+langfun/env/test_utils.py,sha256=Cx4J69r_dLcmDyzEVqOzAZ1Gtk5CnbM2K6nSnRHzIAc,13874
+langfun/env/event_handlers/__init__.py,sha256=sZTUeYpnUMzfH6AghhT-hnBHnFxg5UaH-AOmw3YbyMM,376
+langfun/env/event_handlers/base.py,sha256=w2pdAKfGL-TDTQD2lQ-c_BD3jWeAXF5SYBkAj2GNRz0,7532
+langfun/env/event_handlers/event_logger.py,sha256=4jDV6siGRXR95QFJj_9rnQ_BYMaR6IaPb-wuZ7YzCls,11200
+langfun/env/event_handlers/event_logger_test.py,sha256=OpObIB9IXuuNaRFxOuYcmV4RV6sUTbH5zr3-R7B5UGQ,8536
+langfun-0.1.2.dev202509260805.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202509260805.dist-info/METADATA,sha256=Tee4run8k20LEd7W2ZXxM6E_TzH4g6GFQF-ude-fhS4,7380
+langfun-0.1.2.dev202509260805.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+langfun-0.1.2.dev202509260805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202509260805.dist-info/RECORD,,