haystack

1
import datetime
2
from unittest.mock import Mock, patch
3
import pytest
4

5
from haystack import Pipeline, component
6
from haystack.telemetry._telemetry import pipeline_running
7

8

9
@patch("haystack.telemetry._telemetry.telemetry")
10
def test_pipeline_running(telemetry):
11
    telemetry.send_event = Mock()
12

13
    @component
14
    class Component:
15
        def _get_telemetry_data(self):
16
            return {"key": "values"}
17

18
        @component.output_types(value=int)
19
        def run(self):
20
            pass
21

22
    pipe = Pipeline()
23
    pipe.add_component("component", Component())
24
    pipeline_running(pipe)
25

26
    # First run is always sent
27
    telemetry.send_event.assert_called_once_with(
28
        "Pipeline run (2.x)",
29
        {
30
            "pipeline_id": str(id(pipe)),
31
            "runs": 1,
32
            "components": {"test.test_telemetry.Component": [{"name": "component", "key": "values"}]},
33
        },
34
    )
35

36
    # Running again before one minute has passed should not send another event
37
    telemetry.send_event.reset_mock()
38
    pipeline_running(pipe)
39
    telemetry.send_event.assert_not_called()
40

41
    # Set the last telemetry sent time to pretend one minute has passed
42
    pipe._last_telemetry_sent = pipe._last_telemetry_sent - datetime.timedelta(minutes=1)
43

44
    telemetry.send_event.reset_mock()
45
    pipeline_running(pipe)
46
    telemetry.send_event.assert_called_once_with(
47
        "Pipeline run (2.x)",
48
        {
49
            "pipeline_id": str(id(pipe)),
50
            "runs": 3,
51
            "components": {"test.test_telemetry.Component": [{"name": "component", "key": "values"}]},
52
        },
53
    )
54