Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 8 additions & 6 deletions python/tokenspeed/runtime/utils/env.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,12 +167,14 @@ def override(self, value: Any):
backup_value = os.environ.get(self.name)
backup_set_to_none = self._set_to_none
self.set(value)
yield
if backup_present:
os.environ[self.name] = backup_value
else:
os.environ.pop(self.name, None)
self._set_to_none = backup_set_to_none
try:
yield
finally:
if backup_present:
os.environ[self.name] = backup_value
else:
os.environ.pop(self.name, None)
self._set_to_none = backup_set_to_none

def clear(self):
os.environ.pop(self.name, None)
Expand Down
79 changes: 79 additions & 0 deletions test/runtime/test_env.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Regression tests for scoped runtime environment overrides."""

import os

import pytest

from tokenspeed.runtime.utils.env import EnvStr


class _TestEnvs:
TOKENSPEED_TEST_ENV_FIELD_OVERRIDE = EnvStr("default")


FIELD = _TestEnvs.TOKENSPEED_TEST_ENV_FIELD_OVERRIDE
NAME = "TOKENSPEED_TEST_ENV_FIELD_OVERRIDE"


def test_override_restores_existing_value_after_exception(monkeypatch):
FIELD.clear()
monkeypatch.setenv(NAME, "original")

with pytest.raises(RuntimeError, match="boom"):
with FIELD.override("temporary"):
assert FIELD.get() == "temporary"
raise RuntimeError("boom")

assert os.environ[NAME] == "original"
assert FIELD.get() == "original"
assert not FIELD._set_to_none


def test_override_restores_missing_value_after_exception(monkeypatch):
FIELD.clear()
monkeypatch.delenv(NAME, raising=False)

with pytest.raises(RuntimeError, match="boom"):
with FIELD.override("temporary"):
assert FIELD.get() == "temporary"
raise RuntimeError("boom")

assert NAME not in os.environ
assert FIELD.get() == "default"
assert not FIELD._set_to_none


def test_override_restores_explicit_none_after_exception():
FIELD.clear()
FIELD.set(None)
try:
with pytest.raises(RuntimeError, match="boom"):
with FIELD.override("temporary"):
assert FIELD.get() == "temporary"
raise RuntimeError("boom")

assert NAME not in os.environ
assert FIELD.get() is None
assert FIELD._set_to_none
finally:
FIELD.clear()
Loading