Files
LangBot/tests/unit_tests/plugin/test_worker_mode_quotas.py
T

42 lines
1.9 KiB
Python

from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from langbot.pkg.core.stages.load_config import _complete_runtime_policy_defaults, _apply_env_overrides_to_config
from langbot.pkg.plugin.connector import PluginRuntimeConnector
@pytest.mark.parametrize('mode', ['cloud', 'oss'])
def test_mode_overrides_only_reach_cloud_runtime(mode, monkeypatch):
monkeypatch.setenv('PLUGIN__WORKER__DEDICATED__MAX_CPUS', '0.5')
monkeypatch.setenv('PLUGIN__WORKER__DEDICATED__MAX_MEMORY_MB', '256')
monkeypatch.setenv('PLUGIN__WORKER__SHARED__MAX_CPUS', '1.0')
monkeypatch.setenv('PLUGIN__WORKER__SHARED__MAX_MEMORY_MB', '512')
config = _apply_env_overrides_to_config(_complete_runtime_policy_defaults({'plugin': {'enable': True}}))
app = SimpleNamespace(instance_config=SimpleNamespace(data=config), deployment=SimpleNamespace(mode=mode))
connector = PluginRuntimeConnector(app, AsyncMock())
policy = connector._load_worker_policy()
if mode == 'cloud':
assert policy.dedicated.max_cpus == 0.5
assert policy.dedicated.max_memory_mb == 256
assert policy.shared.max_cpus == 1.0
assert policy.shared.max_memory_mb == 512
else:
assert policy.dedicated is None
assert policy.shared is None
assert policy.max_cpus == 1.0
assert policy.max_memory_mb == 512
def test_cloud_unset_overrides_inherit_existing_limits():
config = _complete_runtime_policy_defaults(
{'plugin': {'enable': True, 'worker': {'max_cpus': 1.5, 'max_memory_mb': 768}}}
)
app = SimpleNamespace(instance_config=SimpleNamespace(data=config), deployment=SimpleNamespace(mode='cloud'))
policy = PluginRuntimeConnector(app, AsyncMock())._load_worker_policy()
for mode in ('dedicated', 'shared-runtime-v1'):
resolved = policy.for_execution_mode(mode, 'shared')
assert resolved.max_cpus == 1.5
assert resolved.max_memory_mb == 768