# 3p
import pymemcache

# project
from ddtrace import Pin
from ddtrace.constants import ANALYTICS_SAMPLE_RATE_KEY
from ddtrace.contrib.pymemcache.patch import patch
from ddtrace.contrib.pymemcache.patch import unpatch
from ddtrace.ext import memcached as memcachedx
from ddtrace.ext import net
from tests.utils import DummyTracer
from tests.utils import TracerTestCase

from .utils import MockSocket


_Client = pymemcache.client.base.Client

TEST_HOST = "localhost"
TEST_PORT = 117711

PYMEMCACHE_VERSION = tuple(int(_) for _ in pymemcache.__version__.split("."))


class PymemcacheClientTestCaseMixin(TracerTestCase):
    """Tests for a patched pymemcache.client.base.Client."""

    def get_spans(self):
        pin = Pin.get_from(pymemcache)
        tracer = pin.tracer
        spans = tracer.pop()
        return spans

    def check_spans(self, num_expected, resources_expected, queries_expected):
        """A helper for validating basic span information."""
        spans = self.get_spans()
        self.assertEqual(num_expected, len(spans))

        for span, resource, query in zip(spans, resources_expected, queries_expected):
            self.assert_is_measured(span)
            self.assertEqual(span.get_tag(net.TARGET_HOST), TEST_HOST)
            self.assertEqual(span.get_metric(net.TARGET_PORT), TEST_PORT)
            self.assertEqual(span.name, memcachedx.CMD)
            self.assertEqual(span.span_type, "cache")
            self.assertEqual(span.service, memcachedx.SERVICE)
            self.assertEqual(span.get_tag(memcachedx.QUERY), query)
            self.assertEqual(span.resource, resource)

        return spans

    def setUp(self):
        patch()

    def tearDown(self):
        unpatch()

    def make_client(self, mock_socket_values, **kwargs):
        tracer = DummyTracer()
        Pin.override(pymemcache, tracer=tracer)
        self.client = pymemcache.client.base.Client((TEST_HOST, TEST_PORT), **kwargs)
        self.client.sock = MockSocket(list(mock_socket_values))
        return self.client

    def test_set_success(self):
        client = self.make_client([b"STORED\r\n"])
        result = client.set(b"key", b"value", noreply=False)
        assert result is True

        self.check_spans(1, ["set"], ["set key"])

    def test_get_many_none_found(self):
        client = self.make_client([b"END\r\n"])
        result = client.get_many([b"key1", b"key2"])
        assert result == {}

        self.check_spans(1, ["get_many"], ["get_many key1 key2"])

    def test_get_multi_none_found(self):
        client = self.make_client([b"END\r\n"])
        result = client.get_multi([b"key1", b"key2"])
        assert result == {}

        self.check_spans(1, ["get_many"], ["get_many key1 key2"])

    def test_delete_not_found(self):
        client = self.make_client([b"NOT_FOUND\r\n"])
        result = client.delete(b"key", noreply=False)
        assert result is False

        self.check_spans(1, ["delete"], ["delete key"])

    def test_incr_found(self):
        client = self.make_client([b"STORED\r\n", b"1\r\n"])
        client.set(b"key", 0, noreply=False)
        result = client.incr(b"key", 1, noreply=False)
        assert result == 1

        self.check_spans(2, ["set", "incr"], ["set key", "incr key"])

    def test_get_found(self):
        client = self.make_client([b"STORED\r\n", b"VALUE key 0 5\r\nvalue\r\nEND\r\n"])
        result = client.set(b"key", b"value", noreply=False)
        result = client.get(b"key")
        assert result == b"value"

        self.check_spans(2, ["set", "get"], ["set key", "get key"])

    def test_decr_found(self):
        client = self.make_client([b"STORED\r\n", b"1\r\n"])
        client.set(b"key", 2, noreply=False)
        result = client.decr(b"key", 1, noreply=False)
        assert result == 1

        self.check_spans(2, ["set", "decr"], ["set key", "decr key"])

    def test_add_stored(self):
        client = self.make_client([b"STORED\r", b"\n"])
        result = client.add(b"key", b"value", noreply=False)
        assert result is True

        self.check_spans(1, ["add"], ["add key"])

    def test_delete_many_found(self):
        client = self.make_client([b"STORED\r", b"\n", b"DELETED\r\n"])
        result = client.add(b"key", b"value", noreply=False)
        result = client.delete_many([b"key"], noreply=False)
        assert result is True

        self.check_spans(2, ["add", "delete_many"], ["add key", "delete_many key"])

    def test_set_many_success(self):
        client = self.make_client([b"STORED\r\n"])
        result = client.set_many({b"key": b"value"}, noreply=False)

        resource = "set_many"
        query = "set_many key"
        if PYMEMCACHE_VERSION[0] == 1:
            assert result is True
        elif PYMEMCACHE_VERSION < (3, 4, 4):
            assert result == []
            if isinstance(client, pymemcache.client.hash.HashClient):
                resource = "set"
                query = "set key"
        else:
            assert result == []

        self.check_spans(1, [resource], [query])

    def test_set_multi_success(self):
        # Should just map to set_many
        client = self.make_client([b"STORED\r\n"])
        result = client.set_multi({b"key": b"value"}, noreply=False)

        resource = "set_many"
        query = "set_many key"
        if PYMEMCACHE_VERSION[0] == 1:
            assert result is True
        elif PYMEMCACHE_VERSION < (3, 4, 4):
            assert result == []
            if isinstance(client, pymemcache.client.hash.HashClient):
                resource = "set"
                query = "set key"
        else:
            assert result == []

        self.check_spans(1, [resource], [query])

    def test_analytics_default(self):
        client = self.make_client([b"STORED\r\n"])
        result = client.set(b"key", b"value", noreply=False)
        assert result is True

        spans = self.get_spans()
        self.assertEqual(len(spans), 1)
        self.assertIsNone(spans[0].get_metric(ANALYTICS_SAMPLE_RATE_KEY))

    def test_analytics_with_rate(self):
        with self.override_config("pymemcache", dict(analytics_enabled=True, analytics_sample_rate=0.5)):
            client = self.make_client([b"STORED\r\n"])
            result = client.set(b"key", b"value", noreply=False)
            assert result is True

        spans = self.get_spans()
        self.assertEqual(len(spans), 1)
        self.assertEqual(spans[0].get_metric(ANALYTICS_SAMPLE_RATE_KEY), 0.5)

    def test_analytics_without_rate(self):
        with self.override_config("pymemcache", dict(analytics_enabled=True)):
            client = self.make_client([b"STORED\r\n"])
            result = client.set(b"key", b"value", noreply=False)
            assert result is True

        spans = self.get_spans()
        self.assertEqual(len(spans), 1)
        self.assertEqual(spans[0].get_metric(ANALYTICS_SAMPLE_RATE_KEY), 1.0)
