OpenKB/tests/conftest.py at 5aeb49636cc95f9fc7ec244da068d9321176d951 · VectifyAI/OpenKB · 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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import inspect
import json
import pytest


@pytest.fixture(autouse=True)
def _reset_extra_headers():
    """Keep the process-wide LLM extra-headers / timeout stashes from leaking across tests."""
    from openkb.config import set_extra_headers, set_timeout
    yield
    set_extra_headers({})
    set_timeout(None)


@pytest.fixture
def close_coroutine_run():
    """Test double for asyncio.run that closes unawaited coroutines."""
    def _run(awaitable):
        if inspect.iscoroutine(awaitable):
            awaitable.close()

    return _run


@pytest.fixture
def raise_after_closing_coroutine(close_coroutine_run):
    def _factory(exc):
        def _run(awaitable):
            close_coroutine_run(awaitable)
            raise exc

        return _run

    return _factory


@pytest.fixture
def kb_dir(tmp_path):
    """Create a minimal knowledge base directory structure for testing."""
    # Raw documents
    (tmp_path / "raw").mkdir()

    # Wiki sub-directories
    (tmp_path / "wiki" / "sources" / "images").mkdir(parents=True)
    (tmp_path / "wiki" / "summaries").mkdir(parents=True)
    (tmp_path / "wiki" / "concepts").mkdir(parents=True)
    (tmp_path / "wiki" / "explorations").mkdir(parents=True)
    (tmp_path / "wiki" / "reports").mkdir(parents=True)

    # .openkb state directory
    openkb_dir = tmp_path / ".openkb"
    openkb_dir.mkdir()

    config_yaml = """\
version: "0.1.0"
embedding_model: text-embedding-3-small
llm_model: gpt-4o-mini
chunk_size: 512
chunk_overlap: 64
"""
    (openkb_dir / "config.yaml").write_text(config_yaml)
    (openkb_dir / "hashes.json").write_text(json.dumps({}))

    return tmp_path


@pytest.fixture
def sample_tree():
    """Return a sample PageIndex tree structure dict for testing."""
    return {
        "doc_name": "Sample Document",
        "doc_description": "A sample document used for unit testing.",
        "structure": [
            {
                "title": "Introduction",
                "node_id": "node-1",
                "start_index": 0,
                "end_index": 120,
                "summary": "Overview of the document topic.",
                "text": "This document introduces the core concepts of the system.",
                "nodes": [
                    {
                        "title": "Background",
                        "node_id": "node-1-1",
                        "start_index": 0,
                        "end_index": 60,
                        "summary": "Historical context.",
                        "text": "Background information on the subject.",
                        "nodes": [],
                    },
                    {
                        "title": "Motivation",
                        "node_id": "node-1-2",
                        "start_index": 61,
                        "end_index": 120,
                        "summary": "Why this work matters.",
                        "text": "Explanation of the motivation behind this work.",
                        "nodes": [],
                    },
                ],
            },
            {
                "title": "Conclusion",
                "node_id": "node-2",
                "start_index": 121,
                "end_index": 200,
                "summary": "Summary of findings.",
                "text": "The system performs well under the described conditions.",
                "nodes": [],
            },
        ],
    }