test_socket.py 文件源码

def test_custom_hostname_resolver(monkeygai):
    class CustomResolver:
        async def getaddrinfo(self, host, port, family, type, proto, flags):
            return ("custom_gai", host, port, family, type, proto, flags)

        async def getnameinfo(self, sockaddr, flags):
            return ("custom_gni", sockaddr, flags)

    cr = CustomResolver()

    assert tsocket.set_custom_hostname_resolver(cr) is None

    # Check that the arguments are all getting passed through.
    # We have to use valid calls to avoid making the underlying system
    # getaddrinfo cranky when it's used for NUMERIC checks.
    for vals in [
        (tsocket.AF_INET, 0, 0, 0),
        (0, tsocket.SOCK_STREAM, 0, 0),
        (0, 0, tsocket.IPPROTO_TCP, 0),
        (0, 0, 0, tsocket.AI_CANONNAME),
    ]:
        assert (
            await tsocket.getaddrinfo(
                "localhost", "foo", *vals
            ) == ("custom_gai", b"localhost", "foo", *vals)
        )

    # IDNA encoding is handled before calling the special object
    got = await tsocket.getaddrinfo("föö", "foo")
    expected = ("custom_gai", b"xn--f-1gaa", "foo", 0, 0, 0, 0)
    assert got == expected

    assert (await tsocket.getnameinfo("a", 0) == ("custom_gni", "a", 0))

    # We can set it back to None
    assert tsocket.set_custom_hostname_resolver(None) is cr

    # And now trio switches back to calling socket.getaddrinfo (specifically
    # our monkeypatched version of socket.getaddrinfo)
    monkeygai.set("x", b"host", "port", family=0, type=0, proto=0, flags=0)
    assert await tsocket.getaddrinfo("host", "port") == "x"