haystack

1
from datetime import datetime
2
from pathlib import Path
3
from typing import Generator
4
from unittest.mock import Mock, patch
5

6
import pytest
7
from openai.types.chat import ChatCompletion, ChatCompletionMessage
8
from openai.types.chat.chat_completion import Choice
9

10
from haystack import tracing
11
from haystack.testing.test_utils import set_all_seeds
12
from test.tracing.utils import SpyingTracer
13

14
set_all_seeds(0)
15

16

17
@pytest.fixture()
18
def mock_tokenizer():
19
    """
20
    Tokenizes the string by splitting on spaces.
21
    """
22
    tokenizer = Mock()
23
    tokenizer.encode = lambda text: text.split()
24
    tokenizer.decode = lambda tokens: " ".join(tokens)
25
    return tokenizer
26

27

28
@pytest.fixture()
29
def test_files_path():
30
    return Path(__file__).parent / "test_files"
31

32

33
@pytest.fixture
34
def mock_chat_completion():
35
    """
36
    Mock the OpenAI API completion response and reuse it for tests
37
    """
38
    with patch("openai.resources.chat.completions.Completions.create") as mock_chat_completion_create:
39
        completion = ChatCompletion(
40
            id="foo",
41
            model="gpt-4",
42
            object="chat.completion",
43
            choices=[
44
                Choice(
45
                    finish_reason="stop",
46
                    logprobs=None,
47
                    index=0,
48
                    message=ChatCompletionMessage(content="Hello world!", role="assistant"),
49
                )
50
            ],
51
            created=int(datetime.now().timestamp()),
52
            usage={"prompt_tokens": 57, "completion_tokens": 40, "total_tokens": 97},
53
        )
54

55
        mock_chat_completion_create.return_value = completion
56
        yield mock_chat_completion_create
57

58

59
@pytest.fixture(autouse=True)
60
def request_blocker(request: pytest.FixtureRequest, monkeypatch):
61
    """
62
    This fixture is applied automatically to all tests.
63
    Those that are not marked as integration will have the requests module
64
    monkeypatched to avoid making HTTP requests by mistake.
65
    """
66
    marker = request.node.get_closest_marker("integration")
67
    if marker is not None:
68
        return
69

70
    def urlopen_mock(self, method, url, *args, **kwargs):
71
        raise RuntimeError(f"The test was about to {method} {self.scheme}://{self.host}{url}")
72

73
    monkeypatch.setattr("urllib3.connectionpool.HTTPConnectionPool.urlopen", urlopen_mock)
74

75

76
@pytest.fixture()
77
def spying_tracer() -> Generator[SpyingTracer, None, None]:
78
    tracer = SpyingTracer()
79
    tracing.enable_tracing(tracer)
80

81
    yield tracer
82

83
    # Make sure to disable tracing after the test to avoid affecting other tests
84
    tracing.disable_tracing()
85