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()