from unittest.mock import MagicMock

from ddtrace.propagation.http import HTTPPropagator

from getstream_connector.utils.extractor import (
    extract_correlation_id,
    extract_trace_context,
)
from tests import constants


def test_extract_trace_context(mocker):
    expected = {"x-datadog-trace-id": "123", "x-datadog-parent-id": "456"}
    message = [{"new": [{constants.TRACE_CONTEXT_KEY: expected}]}]
    mock_extract = mocker.patch.object(HTTPPropagator, "extract", return_value=MagicMock(trace_id=123, span_id=456))

    context = extract_trace_context(message)

    mock_extract.assert_called_once_with(expected)
    assert context.trace_id == 123
    assert context.span_id == 456


def test_extract_trace_context_no_key(mocker):
    message = [{"new": [{}]}]
    mock_warning = mocker.patch("getstream_connector.utils.extractor.logger.warning")

    context = extract_trace_context(message)

    mock_warning.assert_called_once_with("Can't extract trace context from activity.")
    assert context is None


def test_extract_trace_context_invalid_message(mocker):
    message = [{}]
    mock_warning = mocker.patch("getstream_connector.utils.extractor.logger.warning")

    context = extract_trace_context(message)

    mock_warning.assert_called_with("Can't extract trace context from activity.")
    assert context is None


def test_extract_correlation_id():
    message = [{"new": [{constants.CORRELATION_ID_KEY: "corr_id_123"}]}]

    correlation_id = extract_correlation_id(message)
    assert correlation_id == "corr_id_123"


def test_extract_correlation_id_no_key(mocker):
    message = [{"new": [{}]}]
    mock_warning = mocker.patch("getstream_connector.utils.extractor.logger.warning")

    correlation_id = extract_correlation_id(message)

    mock_warning.assert_called_once_with("Can't extract correlation id from activity.")
    assert correlation_id is None


def test_extract_correlation_id_invalid_message(mocker):
    message = [{}]
    mock_warning = mocker.patch("getstream_connector.utils.extractor.logger.warning")

    correlation_id = extract_correlation_id(message)

    mock_warning.assert_called_with("Can't extract correlation id from activity.")
    assert correlation_id is None