[webdriver-client] Refactor wait for transport closed to fix mypy failures.

Differential Revision: https://phabricator.services.mozilla.com/D181909

bugzilla-url: https://bugzilla.mozilla.org/show_bug.cgi?id=1829337
gecko-commit: a1e3851b493a7a6c8baa8d0e9eb917947d46d578
gecko-reviewers: webdriver-reviewers, jgraham, jdescottes
diff --git a/tools/webdriver/webdriver/bidi/modules/session.py b/tools/webdriver/webdriver/bidi/modules/session.py
index c93a9b2..fe1c038 100644
--- a/tools/webdriver/webdriver/bidi/modules/session.py
+++ b/tools/webdriver/webdriver/bidi/modules/session.py
@@ -10,8 +10,8 @@
 
     @end.result
     async def _end(self, result: Mapping[str, Any]) -> Any:
-        if self.session.transport.connection.closed is False:
-            await self.session.transport.connection.wait_closed()
+        if self.session.transport:
+            await self.session.transport.wait_closed()
 
         return result
 
diff --git a/tools/webdriver/webdriver/bidi/transport.py b/tools/webdriver/webdriver/bidi/transport.py
index 2d5d8ee..d61ebad 100644
--- a/tools/webdriver/webdriver/bidi/transport.py
+++ b/tools/webdriver/webdriver/bidi/transport.py
@@ -74,3 +74,7 @@
             if not isinstance(msg, str):
                 raise ValueError("Got a binary message")
             await self.handle(msg)
+
+    async def wait_closed(self) -> None:
+        if self.connection and not self.connection.closed:
+            await self.connection.wait_closed()