haystack/test/others/test_telemetry.py

from pathlib import Path

from unittest.mock import patch, PropertyMock

import pytest

from haystack import telemetry
from haystack.errors import PipelineSchemaError
from haystack.telemetry import (
    NonPrivateParameters,
    send_event,
    enable_writing_events_to_file,
    disable_writing_events_to_file,
    send_custom_event,
    _delete_telemetry_file,
    disable_telemetry,
    enable_telemetry,
    TelemetryFileType,
    _write_telemetry_config,
)


@patch.object(
    NonPrivateParameters, "param_names", return_value=["top_k", "model_name_or_path"], new_callable=PropertyMock
)
def test_private_params_not_tracked(mock_nonprivateparameters):
    params = {"hostname": "private_hostname", "top_k": 2}
    tracked_params = NonPrivateParameters.apply_filter(params)
    expected_params = {"top_k": 2}
    assert tracked_params == expected_params


@patch.object(
    NonPrivateParameters, "param_names", return_value=["top_k", "model_name_or_path"], new_callable=PropertyMock
)
def test_non_private_params_tracked(mock_nonprivateparameters):
    params = {"model_name_or_path": "test-model", "top_k": 2}
    non_private_params = NonPrivateParameters.apply_filter(params)
    assert non_private_params == params


@patch.object(NonPrivateParameters, "param_names", return_value=[], new_callable=PropertyMock)
def test_only_non_private_params(mock_nonprivateparameters):
    non_private_params = NonPrivateParameters.apply_filter({"top_k": 2})
    assert non_private_params == {}


@pytest.mark.integration
@patch("posthog.capture")
@patch.object(
    NonPrivateParameters,
    "param_names",
    return_value=["top_k", "model_name_or_path", "add_isolated_node_eval"],
    new_callable=PropertyMock,
)
# patches are applied in bottom-up order, which is why mock_nonprivateparameters is the first parameter and mock_posthog_capture is the second
def test_send_event_via_decorator(mock_nonprivateparameters, mock_posthog_capture):
    class TestClass:
        @send_event
        def run(self, add_isolated_node_eval: bool = False):
            pass

    test_class = TestClass()
    test_class.run(add_isolated_node_eval=True)
    # todo replace [1] with .kwargs when moving from python 3.7 to 3.8 in CI
    assert mock_posthog_capture.call_args[1]["event"] == "TestClass.run executed"
    assert mock_posthog_capture.call_args[1]["properties"]["add_isolated_node_eval"]


@pytest.mark.integration
@patch("posthog.capture")
def test_send_event_if_custom_error_raised(mock_posthog_capture):
    with pytest.raises(PipelineSchemaError):
        raise PipelineSchemaError
    # todo replace [1] with .kwargs when moving from python 3.7 to 3.8 in CI
    assert mock_posthog_capture.call_args[1]["event"] == "PipelineSchemaError raised"


def num_lines(path: Path):
    if path.is_file():
        with open(path, "r") as f:
            return len(f.readlines())
    return 0


@pytest.mark.integration
@patch("posthog.capture")
def test_write_to_file(mock_posthog_capture, monkeypatch):
    monkeypatch.setattr(telemetry, "LOG_PATH", Path("~/.haystack/telemetry_test.log").expanduser())
    num_lines_before = num_lines(telemetry.LOG_PATH)
    send_custom_event(event="test")
    num_lines_after = num_lines(telemetry.LOG_PATH)
    assert num_lines_before == num_lines_after

    enable_writing_events_to_file()
    num_lines_before = num_lines(telemetry.LOG_PATH)
    send_custom_event(event="test")
    num_lines_after = num_lines(telemetry.LOG_PATH)
    assert num_lines_before + 1 == num_lines_after

    disable_writing_events_to_file()
    num_lines_before = num_lines(telemetry.LOG_PATH)
    send_custom_event(event="test")
    num_lines_after = num_lines(telemetry.LOG_PATH)
    assert num_lines_before == num_lines_after
    _delete_telemetry_file(TelemetryFileType.LOG_FILE)


@pytest.mark.integration
@patch("posthog.capture")
def test_disable_enable_telemetry(mock_posthog_capture, monkeypatch):
    monkeypatch.setattr(telemetry, "HAYSTACK_TELEMETRY_ENABLED", "HAYSTACK_TELEMETRY_ENABLED_TEST")
    monkeypatch.setattr(telemetry, "CONFIG_PATH", Path("~/.haystack/config_test.yaml").expanduser())
    # config_test.yaml doesn't exist yet and won't be created automatically because the global user_id might have been set already by other tests
    _write_telemetry_config()
    send_custom_event(event="test")
    send_custom_event(event="test")
    assert mock_posthog_capture.call_count == 2, "two events should be sent"

    disable_telemetry()
    send_custom_event(event="test")
    assert mock_posthog_capture.call_count == 3, "one additional event should be sent"
    # todo replace [1] with .kwargs when moving from python 3.7 to 3.8 in CI
    assert mock_posthog_capture.call_args[1]["event"] == "telemetry disabled", "a final event should be sent"
    send_custom_event(event="test")
    assert mock_posthog_capture.call_count == 3, "no additional event should be sent"

    enable_telemetry()
    send_custom_event(event="test")
    assert mock_posthog_capture.call_count == 4, "one additional event should be sent"