pytests: refactor to use expect_kresd_close

parent 5b622d2d
......@@ -2,8 +2,6 @@
import time
import pytest
import utils
......@@ -64,11 +62,8 @@ def test_close(kresd_sock):
"""
time.sleep(utils.MAX_TIMEOUT)
with pytest.raises(BrokenPipeError, message="kresd didn't close the connection"):
try:
utils.ping_alive(kresd_sock)
except ConnectionResetError:
pytest.skip('TCP RST')
with utils.expect_kresd_close():
utils.ping_alive(kresd_sock)
def test_slow_lorris_attack(kresd_sock):
......@@ -79,11 +74,8 @@ def test_slow_lorris_attack(kresd_sock):
"""
buff, _ = utils.get_msgbuff()
with pytest.raises(BrokenPipeError, message="kresd didn't close the connection"):
try:
for i in range(len(buff)):
b = buff[i:i+1]
kresd_sock.send(b)
time.sleep(1)
except ConnectionResetError:
pytest.skip('TCP RST')
with utils.expect_kresd_close():
for i in range(len(buff)):
b = buff[i:i+1]
kresd_sock.send(b)
time.sleep(1)
......@@ -16,8 +16,6 @@ TCP_DEFER_ACCEPT, ...), kresd should close the connection.
import time
import pytest
import utils
......@@ -31,12 +29,9 @@ def send_incorrect_repeatedly(sock, buff, delay=1):
"""
end_time = time.time() + utils.MAX_TIMEOUT
with pytest.raises(BrokenPipeError, message="kresd didn't close connection"):
with utils.expect_kresd_close():
while time.time() < end_time:
try:
sock.sendall(buff)
except ConnectionResetError:
pytest.skip("kresd closed connection with TCP RST")
sock.sendall(buff)
time.sleep(delay)
......
from contextlib import contextmanager
import struct
import random
import dns
import dns.message
import pytest
# default net.tcp_in_idle is 10s, TCP_DEFER_ACCEPT 3s, some extra for
......@@ -79,3 +81,12 @@ def ping_alive(sock):
sock.sendall(buff)
answer = receive_parse_answer(sock)
return answer.id == msgid
@contextmanager
def expect_kresd_close():
with pytest.raises(BrokenPipeError, message="kresd didn't close the connection"):
try:
yield
except ConnectionResetError:
pytest.skip("kresd closed connection with TCP RST")
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment