ForgeLM/tests/test_cli_helpers.py at main · HodeTech/ForgeLM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
"""Unit tests for small CLI helper extractions in :mod:`forgelm.cli`.

Covers :func:`forgelm.cli._build_quickstart_inherited_flags` (Phase 10.5
nitpick: argv duplication between train and chat subprocess invocations)
and the module-level GRPO answer regex hoist in :mod:`forgelm.trainer`.
"""

from __future__ import annotations

import argparse
import re

from forgelm.cli import _build_quickstart_inherited_flags
from forgelm.trainer import _ANSWER_PATTERN, _math_reward_fn


def _ns(**kwargs) -> argparse.Namespace:
    return argparse.Namespace(**kwargs)


def test_inherited_flags_default_args_returns_empty_lists() -> None:
    args = _ns()
    train_flags, chat_flags = _build_quickstart_inherited_flags(args)
    assert train_flags == []
    assert chat_flags == []


def test_inherited_flags_propagates_quiet_to_both() -> None:
    args = _ns(quiet=True)
    train_flags, chat_flags = _build_quickstart_inherited_flags(args)
    assert "--quiet" in train_flags
    assert "--quiet" in chat_flags


def test_inherited_flags_propagates_log_level_to_both() -> None:
    args = _ns(log_level="DEBUG")
    train_flags, chat_flags = _build_quickstart_inherited_flags(args)
    # Assert the flag/value pair appears contiguously in both lists.
    assert "--log-level" in train_flags
    assert train_flags[train_flags.index("--log-level") + 1] == "DEBUG"
    assert "--log-level" in chat_flags
    assert chat_flags[chat_flags.index("--log-level") + 1] == "DEBUG"


def test_inherited_flags_propagates_offline_to_both() -> None:
    args = _ns(offline=True)
    train_flags, chat_flags = _build_quickstart_inherited_flags(args)
    assert "--offline" in train_flags
    assert "--offline" in chat_flags


def test_inherited_flags_output_format_not_forwarded() -> None:
    # --output-format json must never be forwarded to the training subprocess.
    # The quickstart parent owns the JSON envelope; forwarding would produce two
    # top-level JSON objects on stdout, making the stream unparseable.
    args = _ns(output_format="json")
    train_flags, chat_flags = _build_quickstart_inherited_flags(args)
    assert "--output-format" not in train_flags
    assert "json" not in train_flags
    assert "--output-format" not in chat_flags
    assert "json" not in chat_flags


def test_inherited_flags_combined() -> None:
    args = _ns(
        quiet=True,
        log_level="INFO",
        offline=True,
        output_format="json",
    )
    train_flags, chat_flags = _build_quickstart_inherited_flags(args)
    assert train_flags == ["--quiet", "--log-level", "INFO", "--offline"]
    assert chat_flags == ["--quiet", "--log-level", "INFO", "--offline"]


def test_math_reward_uses_module_level_pattern() -> None:
    """The hoisted constant must be a compiled regex with case-insensitive flag."""
    assert isinstance(_ANSWER_PATTERN, re.Pattern)
    assert _ANSWER_PATTERN.flags & re.IGNORECASE
    rewards = _math_reward_fn(["Answer: 7"], gold_answer=["7"])
    assert rewards == [1.0]