# This test script was automatically generated by the contrib-patch-tests.py
# script. If you want to make changes to it, you should make sure that you have
# removed the ``_generated`` suffix from the file name, to prevent the content
# from being overwritten by future re-generations.
from ddtrace.contrib.internal.pymongo.patch import get_version
from ddtrace.contrib.internal.pymongo.patch import patch
from ddtrace.contrib.internal.pymongo.patch import pymongo
from ddtrace.contrib.internal.pymongo.patch import unpatch
from tests.contrib.patch import PatchTestCase


_VERSION = pymongo.version_tuple

if _VERSION >= (4, 9):
    from pymongo.synchronous.pool import Connection
    from pymongo.synchronous.server import Server
    from pymongo.synchronous.topology import Topology
elif _VERSION >= (4, 5):
    from pymongo.pool import Connection
    from pymongo.server import Server
    from pymongo.topology import Topology
else:
    from pymongo.pool import SocketInfo as Connection
    from pymongo.server import Server
    from pymongo.topology import Topology


class TestPymongoPatch(PatchTestCase.Base):
    __integration_name__ = "pymongo"
    __module_name__ = "pymongo"
    __patch_func__ = patch
    __unpatch_func__ = unpatch
    __get_version__ = get_version

    def assert_module_patched(self, pymongo):
        self.assert_wrapped(pymongo.MongoClient.__init__)
        self.assert_wrapped(Topology.select_server)

        if _VERSION >= (3, 12):
            self.assert_wrapped(Server.run_operation)
        elif _VERSION >= (3, 9):
            self.assert_wrapped(Server.run_operation_with_response)
        else:
            self.assert_wrapped(Server.send_message_with_response)

        if _VERSION >= (4, 5):
            self.assert_wrapped(Server.checkout)
        else:
            self.assert_wrapped(Server.get_socket)
        self.assert_wrapped(Connection.command)
        self.assert_wrapped(Connection.write_command)

    def assert_not_module_patched(self, pymongo):
        self.assert_not_wrapped(pymongo.MongoClient.__init__)
        self.assert_not_wrapped(Topology.select_server)
        if _VERSION >= (3, 12):
            self.assert_not_wrapped(Server.run_operation)
        elif _VERSION >= (3, 9):
            self.assert_not_wrapped(Server.run_operation_with_response)
        else:
            self.assert_not_wrapped(Server.send_message_with_response)

        if _VERSION >= (4, 5):
            self.assert_not_wrapped(Server.checkout)
        else:
            self.assert_not_wrapped(Server.get_socket)

        self.assert_not_wrapped(Connection.command)
        self.assert_not_wrapped(Connection.write_command)

    def assert_not_module_double_patched(self, pymongo):
        self.assert_not_double_wrapped(pymongo.MongoClient.__init__)
        self.assert_not_double_wrapped(Topology.select_server)
        self.assert_not_double_wrapped(Connection.command)
        self.assert_not_double_wrapped(Connection.write_command)

        if _VERSION >= (3, 12):
            self.assert_not_double_wrapped(Server.run_operation)
        elif _VERSION >= (3, 9):
            self.assert_not_double_wrapped(Server.run_operation_with_response)
        else:
            self.assert_not_double_wrapped(Server.send_message_with_response)

        if _VERSION >= (4, 5):
            self.assert_not_double_wrapped(Server.checkout)
        else:
            self.assert_not_double_wrapped(Server.get_socket)
