feat: add timer

Signed-off-by: Tuan-Dat Tran <tuan-dat.tran@dextradata.com>

programming/tcp-back-to-school/test_main.py

@@ -0,0 +1,51 @@
+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()