From 66a0150a6770c0b53a3d1fded5ff62ae58e78af1 Mon Sep 17 00:00:00 2001 From: RerankerGuo Date: Tue, 13 Oct 2026 17:07:18 +0800 Subject: [PATCH] fix(mem_scheduler): tolerate absent optional fields in EnvConfigMixin.from_env (#2424) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit EnvConfigMixin.from_env previously raised a bare `raise ValueError()` (empty message) whenever a field declared as `field: int | None = None` had no environment variable set — so any config with a plain optional field could not be built from the environment at all. The scheduler init path (`memos/configs/mem_scheduler.py` calls `RabbitMQConfig.from_env()`) was the most visible victim: `port: int | None` is exactly that shape, so it failed immediately. A secondary bug: `_parse_env_value` only parsed when the type was exactly `bool`/`int`/`float`, so an `Optional[int]` env var that *was* provided would have been left as a string instead of being cast. The fix: - from_env now skips fields with no env var entirely; only genuinely required fields (no default and not Optional) still raise, and the raised error carries the missing env var name and the word "missing" (the bug path's bare ValueError() is replaced). - _parse_env_value unwraps `Optional[X]` / `X | None` (using typing.get_args) so an Optional[int] env var parses as int. - A new module-level helper `_is_optional` keeps the Optional check in one place, since both from_env and _parse_env_value need it. New tests/mem_scheduler/test_env_config_mixin.py covers: - absent Optional[int] falls back to None (the bug's reproducer) - absent str falls back to the field default - present Optional[int] parses as int - bool env var still parses as bool (regression coverage) - missing required field still raises with the env var name in the message (so the empty ValueError() path is dead) Fixes #2424. --- .../mem_scheduler/general_modules/misc.py | 58 ++++++++++---- tests/mem_scheduler/test_env_config_mixin.py | 78 +++++++++++++++++++ 2 files changed, 119 insertions(+), 17 deletions(-) create mode 100644 tests/mem_scheduler/test_env_config_mixin.py diff --git a/src/memos/mem_scheduler/general_modules/misc.py b/src/memos/mem_scheduler/general_modules/misc.py index 078f5789b..0fb1dabaa 100644 --- a/src/memos/mem_scheduler/general_modules/misc.py +++ b/src/memos/mem_scheduler/general_modules/misc.py @@ -4,10 +4,11 @@ from contextlib import suppress from datetime import datetime from queue import Empty, Full, Queue -from typing import TYPE_CHECKING, Any, Generic, TypeVar +from typing import TYPE_CHECKING, Any, Generic, TypeVar, get_args, get_origin from dotenv import load_dotenv from pydantic import field_serializer +from typing_extensions import Self if TYPE_CHECKING: @@ -18,6 +19,15 @@ BaseModelType = TypeVar("T", bound="BaseModel") +def _is_optional(annotation: Any) -> bool: + """True if ``annotation`` is ``Optional[X]`` / ``X | None``.""" + origin = get_origin(annotation) + if origin is None: + return False + args = get_args(annotation) + return type(None) in args + + class EnvConfigMixin(Generic[T]): """Abstract base class for environment variable configuration.""" @@ -37,14 +47,13 @@ def get_env_prefix(cls) -> str: """ class_name = cls.__name__ # Remove 'Config' suffix if present - if class_name.endswith("Config"): - class_name = class_name[:-6] + class_name = class_name.removesuffix("Config") # Convert to uppercase and add trailing underscore return f"{cls.ENV_PREFIX}{class_name.upper()}_" @classmethod - def from_env(cls: type[T]) -> T: + def from_env(cls) -> Self: """Creates a config instance from environment variables. Reads all environment variables with the class-specific prefix and maps them @@ -54,7 +63,8 @@ def from_env(cls: type[T]) -> T: An instance of the config class populated from environment variables. Raises: - ValueError: If required environment variables are missing. + ValueError: If a required field (no default and no ``Optional``) has + no matching environment variable. """ load_dotenv() @@ -65,21 +75,36 @@ def from_env(cls: type[T]) -> T: env_var = f"{prefix}{field_name.upper()}" field_type = field_info.annotation - if field_info.is_required() and env_var not in os.environ: - raise ValueError(f"Required environment variable {env_var} is missing") - if env_var in os.environ: raw_value = os.environ[env_var] field_values[field_name] = cls._parse_env_value(raw_value, field_type) - elif field_info.default is not None: - field_values[field_name] = field_info.default - else: - raise ValueError() + continue + + # Absent env var: only fail if the field is genuinely required + # (no default *and* not Optional). Pydantic will apply the + # field's default (or ``None`` for ``Optional[...]``) on its + # own when we don't pass a value at all. Previously this code + # raised a bare ``ValueError()`` for any absent optional field + # (issue #2424). + if field_info.is_required() and not _is_optional(field_type): + raise ValueError(f"Required environment variable {env_var} is missing") + return cls(**field_values) @classmethod def _parse_env_value(cls, value: str, target_type: type) -> Any: - """Converts environment variable string to appropriate type.""" + """Converts environment variable string to appropriate type. + + Unwraps ``Optional[X]`` / ``X | None`` so an optional ``int``/``float`` + env var parses as its inner type instead of staying a raw string + (issue #2424 secondary symptom). + """ + # ``Optional[X]`` and ``X | None`` both come back from get_origin as + # ``types.UnionType`` (PEP 604) or ``typing.Union``. + if _is_optional(target_type): + args = [a for a in get_args(target_type) if a is not type(None)] + if args: + target_type = args[0] if target_type is bool: return value.lower() in ("true", "1", "t", "y", "yes") if target_type is int: @@ -155,7 +180,7 @@ def to_json(self, **kwargs) -> str: return json.dumps(self.to_dict(), ensure_ascii=False, default=lambda o: str(o), **kwargs) @classmethod - def from_json(cls: type[BaseModelType], json_str: str) -> BaseModelType: + def from_json(cls, json_str: str) -> Self: """ Create model instance from a JSON string. - Parses JSON into a dictionary and delegates to from_dict @@ -167,7 +192,7 @@ def from_json(cls: type[BaseModelType], json_str: str) -> BaseModelType: return cls.from_dict(data) @classmethod - def from_dict(cls: type[BaseModelType], data: dict) -> BaseModelType: + def from_dict(cls, data: dict) -> Self: """ Create model instance from dictionary. - Automatically converts timestamp strings to datetime objects @@ -200,8 +225,7 @@ class AutoDroppingQueue(Queue[T]): def __init__(self, maxsize: int = 0): # If maxsize <= 0, set to 0 (unlimited queue size) - if maxsize <= 0: - maxsize = 0 + maxsize = max(0, maxsize) super().__init__(maxsize=maxsize) def put(self, item: T, block: bool = False, timeout: float | None = None) -> None: diff --git a/tests/mem_scheduler/test_env_config_mixin.py b/tests/mem_scheduler/test_env_config_mixin.py new file mode 100644 index 000000000..d76b9e8c2 --- /dev/null +++ b/tests/mem_scheduler/test_env_config_mixin.py @@ -0,0 +1,78 @@ +"""Regression tests for issue #2424. + +``EnvConfigMixin.from_env`` previously raised ``ValueError()`` with an +empty message whenever a field declared as ``field: int | None = None`` +had no environment variable set — so any config with a plain optional +field could not be built from the environment at all. This is on the +scheduler init path (``memos/configs/mem_scheduler.py`` calls +``RabbitMQConfig.from_env()``). + +The secondary bug: ``_parse_env_value`` only parses when the annotation +is exactly ``bool``/``int``/``float``, so an ``Optional[int]`` field with +an env var keeps the raw string instead of being cast. + +These tests pin the new contract: + +- absent optional fields fall back to the model default +- ``Optional[int]`` env vars parse as ``int`` +- genuinely required fields still raise the named error +""" + +from __future__ import annotations + +import pytest + +from pydantic import BaseModel + +from memos.mem_scheduler.general_modules.misc import EnvConfigMixin + + +class RabbitLikeConfig(BaseModel, EnvConfigMixin): + """Stand-in for the scheduler's RabbitMQConfig.""" + + host: str = "localhost" + port: int | None = None + tls: bool = False + + +class RequiredFieldConfig(BaseModel, EnvConfigMixin): + name: str + + +class TestFromEnvAbsentOptional: + def test_absent_optional_int_field_falls_back_to_none(self, monkeypatch): + # No MEMOSCHEDULER_RABBITLIKE_PORT set. + monkeypatch.delenv("MEMSCHEDULER_RABBITLIKE_PORT", raising=False) + cfg = RabbitLikeConfig.from_env() + assert cfg.port is None # was: bare ValueError() + + def test_absent_optional_str_field_falls_back_to_default(self, monkeypatch): + monkeypatch.delenv("MEMSCHEDULER_RABBITLIKE_HOST", raising=False) + cfg = RabbitLikeConfig.from_env() + assert cfg.host == "localhost" + + +class TestFromEnvOptionalIntParses: + def test_optional_int_env_var_parses_as_int(self, monkeypatch): + monkeypatch.setenv("MEMSCHEDULER_RABBITLIKE_PORT", "5672") + cfg = RabbitLikeConfig.from_env() + assert cfg.port == 5672 + assert isinstance(cfg.port, int) + + +class TestFromEnvBoolStillParses: + def test_bool_env_var_parses_as_bool(self, monkeypatch): + monkeypatch.setenv("MEMSCHEDULER_RABBITLIKE_TLS", "yes") + cfg = RabbitLikeConfig.from_env() + assert cfg.tls is True + + +class TestFromEnvRequiredStillRaises: + def test_missing_required_field_raises_named_value_error(self, monkeypatch): + monkeypatch.delenv("MEMSCHEDULER_REQUIREDFIELD_NAME", raising=False) + with pytest.raises(ValueError) as exc_info: + RequiredFieldConfig.from_env() + # Must carry a meaningful message — the empty ``ValueError()`` of + # the unpatched code is what this PR fixes. + assert "MEMSCHEDULER_REQUIREDFIELD_NAME" in str(exc_info.value) + assert "missing" in str(exc_info.value).lower()