import unittest
from mock import Mock, patch

import beeline


class TestRequestsPatch(unittest.TestCase):
    def test_request_fn_injects_headers_and_returns(self):
        from beeline.patch.requests import request  # pylint: disable=bad-option-value,import-outside-toplevel

        with patch('beeline.get_beeline') as m_bl:
            bl = Mock()
            m_bl.return_value = bl

            trace_context = "1;trace_id=foo,parent_id=bar,context=base64value=="

            bl.tracer_impl.http_trace_propagation_hook.return_value = {
                'X-Honeycomb-Trace': trace_context
            }

            # this is the class instance (Session object) that is passed to our request function
            # by wrapt
            m_session = Mock()
            m_session.headers = {}

            # this is our request call that's being wrapped
            m_request = Mock()
            m_request.return_value = Mock(
                headers={'content-type': 'application/json', 'content-length': 23}, status_code=500)
            args = ['get']
            kwargs = {'url': 'http://example.com'}
            ret = request(m_request, m_session, args, kwargs)

            m_request.assert_called_once_with(*args, **kwargs)
            self.assertEqual(ret, m_request.return_value)
            self.assertEqual(
                m_session.headers['X-Honeycomb-Trace'], trace_context)

    def test_request_fn_injects_correct_context(self):
        ''' confirm that the injected trace context references the child span wrapping the request call '''
        from beeline.patch.requests import request  # pylint: disable=bad-option-value,import-outside-toplevel

        _beeline = beeline.Beeline(transmission_impl=Mock())
        # prevent spans from being sent/closed so we can get at the span generated by the request patch
        _beeline.tracer_impl.finish_span = Mock()

        with patch('beeline.get_beeline') as m_bl:
            m_bl.return_value = _beeline

            with _beeline.tracer("parent", trace_id="abc"):
                parent_span = _beeline.tracer_impl.get_active_span()

                # this is the class instance (Session object) that is passed to our request function by wrapt
                m_session = Mock()
                m_session.headers = {}

                # this is our request call that's being wrapped
                m_request = Mock()
                m_request.return_value = Mock(
                    headers={'content-type': 'application/json', 'content-length': 23}, status_code=500)
                args = ['get']
                kwargs = {'url': 'http://example.com'}
                ret = request(m_request, m_session, args, kwargs)

                m_request.assert_called_once_with(*args, **kwargs)
                self.assertEqual(ret, m_request.return_value)

                self.assertEqual(len(_beeline.tracer_impl._trace.stack), 2)
                child_span = _beeline.tracer_impl.get_active_span()
                # we should have two distinct spans in this trace
                self.assertNotEqual(child_span.id, parent_span.id)
                self.assertEqual(child_span.parent_id, parent_span.id)

                trace_context = m_session.headers['X-Honeycomb-Trace']
                trace_id, parent_id, _ = beeline.trace.unmarshal_trace_context(
                    trace_context)
                # confirm the trace context parent is the child span, not the parent span
                self.assertEqual(child_span.trace_id, trace_id)
                self.assertEqual(child_span.id, parent_id)
                # should be the same trace
                self.assertEqual(parent_span.trace_id, trace_id)

    def test_request_fn_works_without_init(self):
        ''' verify that the requests wrapper works even if the beeline is not initialized '''
        from beeline.patch.requests import request  # pylint: disable=bad-option-value,import-outside-toplevel

        # this is the class instance (Session object) that is passed to our request function by wrapt
        m_session = Mock()
        m_session.headers = {}

        m_request = Mock()
        m_request.return_value = Mock(
            headers={'content-type': 'application/json', 'content-length': 23}, status_code=500)
        args = ['get']
        kwargs = {'url': 'http://example.com'}
        ret = request(m_request, m_session, args, kwargs)

        m_request.assert_called_once_with(*args, **kwargs)
        self.assertEqual(ret, m_request.return_value)
