import sqlalchemy

from ddtrace import Pin
from ddtrace.constants import ANALYTICS_SAMPLE_RATE_KEY
from ddtrace.contrib.sqlalchemy import patch
from ddtrace.contrib.sqlalchemy import unpatch
from tests.utils import TracerTestCase
from tests.utils import assert_is_measured

from ..config import POSTGRES_CONFIG


class SQLAlchemyPatchTestCase(TracerTestCase):
    """TestCase that checks if the engine is properly traced
    when the `patch()` method is used.
    """

    def setUp(self):
        super(SQLAlchemyPatchTestCase, self).setUp()

        # create a traced engine with the given arguments
        # and configure the current PIN instance
        patch()
        dsn = "postgresql://%(user)s:%(password)s@%(host)s:%(port)s/%(dbname)s" % POSTGRES_CONFIG
        self.engine = sqlalchemy.create_engine(dsn)
        Pin.override(self.engine, tracer=self.tracer)

        # prepare a connection
        self.conn = self.engine.connect()

    def tearDown(self):
        super(SQLAlchemyPatchTestCase, self).tearDown()

        # clear the database and dispose the engine
        self.conn.close()
        self.engine.dispose()
        unpatch()

    def test_engine_traced(self):
        # ensures that the engine is traced
        rows = self.conn.execute("SELECT 1").fetchall()
        assert len(rows) == 1

        traces = self.pop_traces()
        # trace composition
        assert len(traces) == 1
        assert len(traces[0]) == 1
        span = traces[0][0]
        # check subset of span fields
        assert_is_measured(span)
        assert span.name == "postgres.query"
        assert span.service == "postgres"
        assert span.error == 0
        assert span.duration > 0

    def test_engine_pin_service(self):
        # ensures that the engine service is updated with the PIN object
        Pin.override(self.engine, service="replica-db")
        rows = self.conn.execute("SELECT 1").fetchall()
        assert len(rows) == 1

        traces = self.pop_traces()
        # trace composition
        assert len(traces) == 1
        assert len(traces[0]) == 1
        span = traces[0][0]
        # check subset of span fields
        assert_is_measured(span)
        assert span.name == "postgres.query"
        assert span.service == "replica-db"
        assert span.error == 0
        assert span.duration > 0

    def test_analytics_sample_rate(self):
        # [ <config>, <analytics sample rate metric value> ]
        matrix = [
            # Default, not enabled, not set
            [dict(), None],
            # Not enabled, but sample rate set
            [dict(analytics_sample_rate=0.5), None],
            # Enabled and rate set
            [dict(analytics_enabled=True, analytics_sample_rate=0.5), 0.5],
            [dict(analytics_enabled=True, analytics_sample_rate=1), 1.0],
            [dict(analytics_enabled=True, analytics_sample_rate=0), 0],
            [dict(analytics_enabled=True, analytics_sample_rate=True), 1.0],
            [dict(analytics_enabled=True, analytics_sample_rate=False), 0],
            # Disabled and rate set
            [dict(analytics_enabled=False, analytics_sample_rate=0.5), None],
            # Enabled and rate not set
            [dict(analytics_enabled=True), 1.0],
        ]
        for config, metric_value in matrix:
            with self.override_config("sqlalchemy", config):
                self.conn.execute("SELECT 1").fetchall()

                root = self.get_root_span()
                assert_is_measured(root)
                root.assert_matches(name="postgres.query")

                # If the value is None assert it was not set, otherwise assert the expected value
                # DEV: root.assert_metrics(metrics, exact=True) won't work here since we have another sample
                #      rate keys getting added
                if metric_value is None:
                    assert ANALYTICS_SAMPLE_RATE_KEY not in root.metrics
                else:
                    assert root.metrics[ANALYTICS_SAMPLE_RATE_KEY] == metric_value
                self.reset()
