from pathlib import Path
import unittest


ROOT = Path(__file__).resolve().parents[1]


class WebsocketProxyClientLifecycleTests(unittest.TestCase):
    def test_each_proxy_connection_has_its_own_cleanup_key(self) -> None:
        source = (ROOT / "ws" / "main.py").read_text(encoding="utf-8")

        self.assertIn("client_key = id(websocket)", source)
        self.assertNotIn("if client_address in self.__clients", source)
        self.assertIn("client = self.__clients.get(client_key)", source)
        self.assertIn("self.__clients.pop(client_key, None)", source)

    def test_cleanup_absorbs_a_reader_closed_by_the_browser(self) -> None:
        source = (ROOT / "ws" / "main.py").read_text(encoding="utf-8")
        cleanup = source[source.index("client = self.__clients.get(client_key)"):]

        self.assertIn("except websockets.exceptions.ConnectionClosed:", cleanup)


if __name__ == "__main__":
    unittest.main()