/
testing.py
64 lines (50 loc) · 1.91 KB
/
testing.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
"""Utilities for defining Python tests."""
import socket
from platform import system
from typing import Any, TypedDict
PytestSkip = TypedDict("PytestSkip", {"condition": bool, "reason": str})
def has_ipv6() -> bool:
"""Check whether IPv6 is enabled on this host."""
# connection error in macos, still need some fixes.
if system() not in ("Linux", "Windows"):
return False
if socket.has_ipv6:
try:
with socket.socket(
socket.AF_INET6, socket.SOCK_STREAM
) as server, socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as client:
server.bind(("::1", 0))
port = server.getsockname()[1]
server.listen()
client.connect(("::1", port))
conn, _ = server.accept()
client.sendall("abc".encode())
msg = conn.recv(3).decode()
# if the code can be executed to this point, the message should be
# correct.
assert msg == "abc"
return True
except OSError:
pass
return False
def skip_ipv6() -> PytestSkip:
"""PyTest skip mark for IPv6."""
return {"condition": not has_ipv6(), "reason": "IPv6 is required to be enabled."}
def timeout(sec: int, *args: Any, enable: bool = False, **kwargs: Any) -> Any:
"""Make a pytest mark for the `pytest-timeout` package.
Parameters
----------
sec :
Timeout seconds.
enable :
Control whether timeout should be applied, used for debugging.
Returns
-------
pytest.mark.timeout
"""
import pytest # pylint: disable=import-error
# This is disabled for now due to regression caused by conflicts between federated
# learning build and the CI container environment.
if enable:
return pytest.mark.timeout(sec, *args, **kwargs)
return pytest.mark.timeout(None, *args, **kwargs)