Repository navigation
Expand file tree
/
Copy pathconftest.py
More file actions
171 lines (140 loc) · 5.05 KB
/
Copy pathconftest.py
File metadata and controls
171 lines (140 loc) · 5.05 KB
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
import os
# import socket
from typing import Any, List, Optional
import openai
import pytest
from llama_index.core.llms.types import LLMMetadata
from llama_index.llm_predictor.base import LLMPredictor
from llama_index.llms.mock import MockLLM
from llama_index.node_parser.text import SentenceSplitter, TokenTextSplitter
from llama_index.service_context import ServiceContext
from tests.indices.vector_store.mock_services import MockEmbedding
from tests.mock_utils.mock_predict import (
patch_llmpredictor_apredict,
patch_llmpredictor_predict,
)
from tests.mock_utils.mock_text_splitter import patch_token_splitter_newline
# @pytest.fixture(autouse=True)
# def no_networking(monkeypatch: pytest.MonkeyPatch) -> None:
# def deny_network(*args: Any, **kwargs: Any) -> None:
# raise RuntimeError("Network access denied for test")
# monkeypatch.setattr(socket, "socket", deny_network)
@pytest.fixture()
def allow_networking(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.undo()
@pytest.fixture()
def patch_token_text_splitter(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(SentenceSplitter, "split_text", patch_token_splitter_newline)
monkeypatch.setattr(
SentenceSplitter,
"split_text_metadata_aware",
patch_token_splitter_newline,
)
monkeypatch.setattr(TokenTextSplitter, "split_text", patch_token_splitter_newline)
monkeypatch.setattr(
TokenTextSplitter, "split_text_metadata_aware", patch_token_splitter_newline
)
@pytest.fixture()
def patch_llm_predictor(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(
LLMPredictor,
"predict",
patch_llmpredictor_predict,
)
monkeypatch.setattr(
LLMPredictor,
"apredict",
patch_llmpredictor_apredict,
)
monkeypatch.setattr(
LLMPredictor,
"llm",
MockLLM(),
)
monkeypatch.setattr(
LLMPredictor,
"metadata",
LLMMetadata(),
)
monkeypatch.setattr(
MockLLM,
"predict",
patch_llmpredictor_predict,
)
monkeypatch.setattr(
MockLLM,
"apredict",
patch_llmpredictor_apredict,
)
monkeypatch.setattr(
MockLLM,
"metadata",
LLMMetadata(),
)
@pytest.fixture()
def mock_service_context(
patch_token_text_splitter: Any,
patch_llm_predictor: Any,
) -> ServiceContext:
return ServiceContext.from_defaults(embed_model=MockEmbedding())
@pytest.fixture()
def mock_llm() -> MockLLM:
return MockLLM()
@pytest.fixture(autouse=True)
def mock_openai_credentials() -> None:
if not os.environ.get("OPENAI_API_KEY"):
os.environ["OPENAI_API_KEY"] = "sk-" + ("a" * 48)
class CachedOpenAIApiKeys:
"""
Saves the users' OpenAI API key and OpenAI API type either in
the environment variable or set to the library itself.
This allows us to run tests by setting it without plowing over
the local environment.
"""
def __init__(
self,
set_env_key_to: Optional[str] = "",
set_library_key_to: Optional[str] = None,
set_fake_key: bool = False,
set_env_type_to: Optional[str] = "",
set_library_type_to: str = "open_ai", # default value in openai package
):
self.set_env_key_to = set_env_key_to
self.set_library_key_to = set_library_key_to
self.set_fake_key = set_fake_key
self.set_env_type_to = set_env_type_to
self.set_library_type_to = set_library_type_to
def __enter__(self) -> None:
self.api_env_variable_was = os.environ.get("OPENAI_API_KEY", "")
self.api_env_type_was = os.environ.get("OPENAI_API_TYPE", "")
self.openai_api_key_was = openai.api_key
self.openai_api_type_was = openai.api_type
os.environ["OPENAI_API_KEY"] = str(self.set_env_key_to)
os.environ["OPENAI_API_TYPE"] = str(self.set_env_type_to)
if self.set_fake_key:
os.environ["OPENAI_API_KEY"] = "sk-" + "a" * 48
# No matter what, set the environment variable back to what it was
def __exit__(self, *exc: object) -> None:
os.environ["OPENAI_API_KEY"] = str(self.api_env_variable_was)
os.environ["OPENAI_API_TYPE"] = str(self.api_env_type_was)
openai.api_key = self.openai_api_key_was
openai.api_type = self.openai_api_type_was
def pytest_addoption(parser: pytest.Parser) -> None:
parser.addoption(
"--integration",
action="store_true",
default=False,
help="run integration tests",
)
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line("markers", "integration: mark test as integration")
def pytest_collection_modifyitems(
config: pytest.Config, items: List[pytest.Item]
) -> None:
if config.getoption("--integration"):
# --integration given in cli: do not skip integration tests
return
skip_integration = pytest.mark.skip(reason="need --integration option to run")
for item in items:
if "integration" in item.keywords:
item.add_marker(skip_integration)