CodeWiki/tests/test_model_flags.py at df96a09585ce27b7201b59cbb30c1d2f8bdeaedb · FSoft-AI4Code/CodeWiki · 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
"""`codewiki generate --main-model/--cluster-model/--fallback-model` override the saved config."""

from types import SimpleNamespace

import pytest
from click.testing import CliRunner

import codewiki.cli.commands.generate as gen


class Stop(Exception):
    pass


@pytest.fixture
def captured(tmp_path, monkeypatch):
    seen = {}
    cfg = SimpleNamespace(
        main_model="saved-main",
        cluster_model="saved-cluster",
        fallback_model="saved-fallback",
        base_url="http://x",
        provider="openai-compatible",
        aws_region="us-east-1",
        agent_instructions=None,
        max_tokens=1,
        max_token_per_module=1,
        max_token_per_leaf_module=1,
        max_depth=1,
        request_limit=1,
        agent_retries=1,
        use_gitignore=True,
        prompt_caching=False,
    )

    class FakeManager:
        def load(self):
            return True

        def is_configured(self):
            return True

        def get_config(self):
            return cfg

        def get_api_key(self):
            return "k"

    def fake_generator(**kwargs):
        seen.update(kwargs["config"])
        raise Stop

    monkeypatch.setattr(gen, "ConfigManager", FakeManager)
    monkeypatch.setattr(gen, "CLIDocumentationGenerator", fake_generator)
    (tmp_path / "a.py").write_text("def f():\n    return 1\n")
    monkeypatch.chdir(tmp_path)
    return seen


def test_saved_models_by_default(captured):
    CliRunner().invoke(gen.generate_command, ["-o", "out"])
    assert (captured["main_model"], captured["cluster_model"], captured["fallback_model"]) == (
        "saved-main",
        "saved-cluster",
        "saved-fallback",
    )


def test_flags_override_saved_models(captured):
    CliRunner().invoke(
        gen.generate_command,
        ["-o", "out", "--main-model", "m", "--cluster-model", "c", "--fallback-model", "f"],
    )
    assert (captured["main_model"], captured["cluster_model"], captured["fallback_model"]) == (
        "m",
        "c",
        "f",
    )