"""Tests for trace utils module."""

import base64
import json

from notifications_delivery.utils import trace_utils


def test_extract_trace_context(mocker):
    """Test extract_trace_context."""
    mock_extract_tc = mocker.patch.object(
        trace_utils.extractor, 'extract_trace_context', return_value='TRACE_CONTEXT'
    )
    mock_extract_cid = mocker.patch.object(
        trace_utils.extractor, 'extract_correlation_id', return_value='CORRELATION_ID'
    )

    raw = [
        {
            'new': [{
                'data': 'some'
            }]
        }
    ]
    message = type('Message', (object,), {'body': base64.b64encode(json.dumps(raw).encode())})()

    trace_context, correlation_id = trace_utils.extract_trace_context(message)
    assert trace_context == 'TRACE_CONTEXT'
    assert correlation_id == 'CORRELATION_ID'
    mock_extract_tc.assert_called_once_with(raw)
    mock_extract_cid.assert_called_once_with(raw)


def test_extract_trace_context_not_base64(mocker):
    """Test extract_trace_context when message is not base64 encoded."""
    mock_extract_tc = mocker.patch.object(
        trace_utils.extractor, 'extract_trace_context', return_value='TRACE_CONTEXT'
    )
    mock_extract_cid = mocker.patch.object(
        trace_utils.extractor, 'extract_correlation_id', return_value='CORRELATION_ID'
    )

    raw = [
        {
            'new': [{
                'data': 'some'
            }]
        }
    ]
    message = type('Message', (object,), {'body': json.dumps(raw).encode()})()

    trace_context, correlation_id = trace_utils.extract_trace_context(message)
    assert trace_context == 'TRACE_CONTEXT'
    assert correlation_id == 'CORRELATION_ID'
    mock_extract_tc.assert_called_once_with(raw)
    mock_extract_cid.assert_called_once_with(raw)


def test_extract_trace_context_not_json(mocker):
    """Test extract_trace_context when message is not json encoded."""
    mock_extract_tc = mocker.patch.object(
        trace_utils.extractor, 'extract_trace_context', return_value='TRACE_CONTEXT'
    )
    mock_extract_cid = mocker.patch.object(
        trace_utils.extractor, 'extract_correlation_id', return_value='CORRELATION_ID'
    )

    raw = 'data'
    message = type('Message', (object,), {'body': raw})()

    trace_context, correlation_id = trace_utils.extract_trace_context(message)
    assert trace_context is None
    assert correlation_id is None
    mock_extract_tc.assert_not_called()
    mock_extract_cid.assert_not_called()