test_pymemcache.py

# Copyright The OpenTelemetry Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest import mock

import pymemcache
from pymemcache.exceptions import (
    MemcacheClientError,
    MemcacheIllegalInputError,
    MemcacheServerError,
    MemcacheUnknownCommandError,
    MemcacheUnknownError,
)

from opentelemetry import trace as trace_api
from opentelemetry.instrumentation.pymemcache import PymemcacheInstrumentor
from opentelemetry.test.test_base import TestBase
from opentelemetry.trace import get_tracer

from .utils import MockSocket, _str

TEST_HOST = "localhost"
TEST_PORT = 117711


class PymemcacheClientTestCase(
    TestBase
):  # pylint: disable=too-many-public-methods
    """ Tests for a patched pymemcache.client.base.Client. """

    def setUp(self):
        super().setUp()
        PymemcacheInstrumentor().instrument()

        # pylint: disable=protected-access
        self.tracer = get_tracer(__name__)

    def tearDown(self):
        super().tearDown()
        PymemcacheInstrumentor().uninstrument()

    def make_client(self, mock_socket_values, **kwargs):
        # pylint: disable=attribute-defined-outside-init
        self.client = pymemcache.client.base.Client(
            (TEST_HOST, TEST_PORT), **kwargs
        )
        self.client.sock = MockSocket(list(mock_socket_values))
        return self.client

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

        for span, query in zip(spans, queries_expected):
            command, *_ = query.split(" ")
            self.assertEqual(span.name, command)
            self.assertIs(span.kind, trace_api.SpanKind.CLIENT)
            self.assertEqual(span.attributes["net.peer.name"], TEST_HOST)
            self.assertEqual(span.attributes["net.peer.port"], TEST_PORT)
            self.assertEqual(span.attributes["db.system"], "memcached")
            self.assertEqual(span.attributes["db.statement"], query)

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

        spans = self.memory_exporter.get_finished_spans()

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

    def test_set_not_recording(self):
        mock_tracer = mock.Mock()
        mock_span = mock.Mock()
        mock_span.is_recording.return_value = False
        mock_tracer.start_span.return_value = mock_span
        mock_tracer.use_span.return_value.__enter__ = mock_span
        mock_tracer.use_span.return_value.__exit__ = True
        with mock.patch("opentelemetry.trace.get_tracer") as tracer:
            tracer.return_value = mock_tracer
            client = self.make_client([b"STORED\r\n"])
            result = client.set(b"key", b"value", noreply=False)
            self.assertTrue(result)
            self.assertFalse(mock_span.is_recording())
            self.assertTrue(mock_span.is_recording.called)
            self.assertFalse(mock_span.set_attribute.called)
            self.assertFalse(mock_span.set_status.called)

    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 == {}

        spans = self.memory_exporter.get_finished_spans()

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

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

        spans = self.memory_exporter.get_finished_spans()

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

    def test_set_multi_success(self):
        client = self.make_client([b"STORED\r\n"])
        # Alias for set_many, a convienance function that calls set for every key
        result = client.set_multi({b"key": b"value"}, noreply=False)
        self.assertTrue(result)

        spans = self.memory_exporter.get_finished_spans()

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

    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

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["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

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 2, ["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"

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 2, ["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

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 2, ["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)
        self.assertTrue(result)

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["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)
        # a convienance function that calls delete for every key
        result = client.delete_many([b"key"], noreply=False)
        self.assertTrue(result)

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(
            spans, 3, ["add key", "delete key", "delete_many key"]
        )

    def test_set_many_success(self):
        client = self.make_client([b"STORED\r\n"])
        # a convienance function that calls set for every key
        result = client.set_many({b"key": b"value"}, noreply=False)
        self.assertTrue(result)

        spans = self.memory_exporter.get_finished_spans()

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

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

        spans = self.memory_exporter.get_finished_spans()

        self.assertEqual(len(spans), 2)
        self.assertEqual(spans[0].attributes["net.peer.name"], TEST_HOST)
        self.assertEqual(spans[0].attributes["net.peer.port"], TEST_PORT)

    def test_append_stored(self):
        client = self.make_client([b"STORED\r\n"])
        result = client.append(b"key", b"value", noreply=False)
        self.assertTrue(result)

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["append key"])

    def test_prepend_stored(self):
        client = self.make_client([b"STORED\r\n"])
        result = client.prepend(b"key", b"value", noreply=False)
        self.assertTrue(result)

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["prepend key"])

    def test_cas_stored(self):
        client = self.make_client([b"STORED\r\n"])
        result = client.cas(b"key", b"value", b"cas", noreply=False)
        self.assertTrue(result)

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["cas key"])

    def test_cas_exists(self):
        client = self.make_client([b"EXISTS\r\n"])
        result = client.cas(b"key", b"value", b"cas", noreply=False)
        assert result is False

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["cas key"])

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

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["cas key"])

    def test_delete_exception(self):
        client = self.make_client([Exception("fail")])

        def _delete():
            client.delete(b"key", noreply=False)

        with self.assertRaises(Exception):
            _delete()

        spans = self.memory_exporter.get_finished_spans()

        span = spans[0]

        self.assertFalse(span.status.is_ok)

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

    def test_flush_all(self):
        client = self.make_client([b"OK\r\n"])
        result = client.flush_all(noreply=False)
        self.assertTrue(result)

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["flush_all"])

    def test_incr_exception(self):
        client = self.make_client([Exception("fail")])

        def _incr():
            client.incr(b"key", 1)

        with self.assertRaises(Exception):
            _incr()

        spans = self.memory_exporter.get_finished_spans()

        span = spans[0]

        self.assertFalse(span.status.is_ok)

        self.check_spans(spans, 1, ["incr key"])

    def test_get_error(self):
        client = self.make_client([b"ERROR\r\n"])

        def _get():
            client.get(b"key")

        with self.assertRaises(MemcacheUnknownCommandError):
            _get()

        spans = self.memory_exporter.get_finished_spans()

        span = spans[0]

        self.assertFalse(span.status.is_ok)

        self.check_spans(spans, 1, ["get key"])

    def test_get_unknown_error(self):
        client = self.make_client([b"foobarbaz\r\n"])

        def _get():
            client.get(b"key")

        with self.assertRaises(MemcacheUnknownError):
            _get()

        spans = self.memory_exporter.get_finished_spans()

        span = spans[0]

        self.assertFalse(span.status.is_ok)

        self.check_spans(spans, 1, ["get key"])

    def test_gets_found(self):
        client = self.make_client([b"VALUE key 0 5 10\r\nvalue\r\nEND\r\n"])
        result = client.gets(b"key")
        assert result == (b"value", b"10")

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["gets key"])

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

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["touch key"])

    def test_set_client_error(self):
        client = self.make_client([b"CLIENT_ERROR some message\r\n"])

        def _set():
            client.set("key", "value", noreply=False)

        with self.assertRaises(MemcacheClientError):
            _set()

        spans = self.memory_exporter.get_finished_spans()

        span = spans[0]

        self.assertFalse(span.status.is_ok)

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

    def test_set_server_error(self):
        client = self.make_client([b"SERVER_ERROR some message\r\n"])

        def _set():
            client.set(b"key", b"value", noreply=False)

        with self.assertRaises(MemcacheServerError):
            _set()

        spans = self.memory_exporter.get_finished_spans()

        span = spans[0]

        self.assertFalse(span.status.is_ok)

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

    def test_set_key_with_space(self):
        client = self.make_client([b""])

        def _set():
            client.set(b"key has space", b"value", noreply=False)

        with self.assertRaises(MemcacheIllegalInputError):
            _set()

        spans = self.memory_exporter.get_finished_spans()

        span = spans[0]

        self.assertFalse(span.status.is_ok)

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

    def test_quit(self):
        client = self.make_client([])
        assert client.quit() is None

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["quit"])

    def test_replace_not_stored(self):
        client = self.make_client([b"NOT_STORED\r\n"])
        result = client.replace(b"key", b"value", noreply=False)
        assert result is False

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["replace key"])

    def test_version_success(self):
        client = self.make_client(
            [b"VERSION 1.2.3\r\n"], default_noreply=False
        )
        result = client.version()
        assert result == b"1.2.3"

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["version"])

    def test_stats(self):
        client = self.make_client([b"STAT fake_stats 1\r\n", b"END\r\n"])
        result = client.stats()
        assert client.sock.send_bufs == [b"stats \r\n"]
        assert result == {b"fake_stats": 1}

        spans = self.memory_exporter.get_finished_spans()

        self.check_spans(spans, 1, ["stats"])

    def test_uninstrumented(self):
        PymemcacheInstrumentor().uninstrument()

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

        spans = self.memory_exporter.get_finished_spans()

        self.assertEqual(len(spans), 0)

        PymemcacheInstrumentor().instrument()


class PymemcacheHashClientTestCase(TestBase):
    """ Tests for a patched pymemcache.client.hash.HashClient. """

    def setUp(self):
        super().setUp()
        PymemcacheInstrumentor().instrument()

        # pylint: disable=protected-access
        self.tracer = get_tracer(__name__)

    def tearDown(self):
        super().tearDown()
        PymemcacheInstrumentor().uninstrument()

    def make_client_pool(
        self, hostname, mock_socket_values, serializer=None, **kwargs
    ):  # pylint: disable=no-self-use
        mock_client = pymemcache.client.base.Client(
            hostname, serializer=serializer, **kwargs
        )
        mock_client.sock = MockSocket(mock_socket_values)
        client = pymemcache.client.base.PooledClient(
            hostname, serializer=serializer
        )
        client.client_pool = pymemcache.pool.ObjectPool(lambda: mock_client)
        return mock_client

    def make_client(self, *mock_socket_values, **kwargs):
        current_port = TEST_PORT

        # pylint: disable=import-outside-toplevel
        from pymemcache.client.hash import HashClient

        # pylint: disable=attribute-defined-outside-init
        self.client = HashClient([], **kwargs)
        ip = TEST_HOST

        for vals in mock_socket_values:
            url_string = "{}:{}".format(ip, current_port)
            clnt_pool = self.make_client_pool(
                (ip, current_port), vals, **kwargs
            )
            self.client.clients[url_string] = clnt_pool
            self.client.hasher.add_node(url_string)
            current_port += 1
        return self.client

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

        for span, query in zip(spans, queries_expected):
            command, *_ = query.split(" ")
            self.assertEqual(span.name, command)
            self.assertIs(span.kind, trace_api.SpanKind.CLIENT)
            self.assertEqual(span.attributes["net.peer.name"], TEST_HOST)
            self.assertEqual(span.attributes["net.peer.port"], TEST_PORT)
            self.assertEqual(span.attributes["db.system"], "memcached")
            self.assertEqual(span.attributes["db.statement"], query)

    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)
        self.assertTrue(result)

        spans = self.memory_exporter.get_finished_spans()

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