52 lines
1.2 KiB
Python
52 lines
1.2 KiB
Python
import io
|
|
import unittest
|
|
from contextlib import redirect_stdout
|
|
from unittest.mock import patch
|
|
|
|
import main
|
|
|
|
|
|
class FakeConn:
|
|
def __init__(self) -> None:
|
|
self.closed = False
|
|
self.sent = []
|
|
self._responses = [
|
|
b"Calculate the square root of 4 and multiply by 5 = ",
|
|
b"[+] Good job ! Here is your flag: RM{test}",
|
|
]
|
|
|
|
def recv(self, _size: int) -> bytes:
|
|
return self._responses.pop(0)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, _exc_type, _exc, _tb) -> None:
|
|
self.close()
|
|
|
|
def sendall(self, payload: bytes) -> None:
|
|
self.sent.append(payload)
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class MainTests(unittest.TestCase):
|
|
def test_main_prints_connection_timer_after_closing(self) -> None:
|
|
conn = FakeConn()
|
|
out = io.StringIO()
|
|
|
|
with (
|
|
patch("main.socket.create_connection", return_value=conn),
|
|
redirect_stdout(out),
|
|
):
|
|
main.main()
|
|
|
|
output = out.getvalue()
|
|
self.assertIn("Connection lifetime:", output)
|
|
self.assertTrue(conn.closed)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|