Skip to content
Merged
16 changes: 15 additions & 1 deletion playwright/_impl/_browser_type.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
from pathlib import Path
from typing import Dict, List, Optional, Union, cast

Expand Down Expand Up @@ -179,11 +180,24 @@ async def connect(
self._connection._object_factory,
transport,
)
await connection._transport.start()
connection._is_sync = self._connection._is_sync
connection._loop = self._connection._loop
connection._loop.create_task(connection.run())
obj = asyncio.create_task(
connection.wait_for_object_with_known_name("Playwright")
)
done, pending = await asyncio.wait(
{
obj,
connection._transport.on_error_future, # type: ignore
},
return_when=asyncio.FIRST_COMPLETED,
)
if not obj.done():
obj.cancel()
playwright = next(iter(done)).result()
self._connection._child_ws_connections.append(connection)
playwright = await connection.wait_for_object_with_known_name("Playwright")
pre_launched_browser = playwright._initializer.get("preLaunchedBrowser")
assert pre_launched_browser
browser = cast(Browser, from_channel(pre_launched_browser))
Expand Down
50 changes: 37 additions & 13 deletions playwright/_impl/_transport.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ def _get_stderr_fileno() -> Optional[int]:

class Transport(ABC):
def __init__(self) -> None:
self.on_error_future: asyncio.Future
self.on_message = lambda _: None

@abstractmethod
Expand All @@ -55,9 +56,14 @@ def dispose(self) -> None:
async def wait_until_stopped(self) -> None:
pass

async def run(self) -> None:
async def start(self) -> None:
if not hasattr(self, "on_error_future"):
self.on_error_future = asyncio.Future()
self._loop = asyncio.get_running_loop()
self.on_error_future: asyncio.Future = asyncio.Future()

@abstractmethod
async def run(self) -> None:
pass

@abstractmethod
def send(self, message: Dict) -> None:
Expand Down Expand Up @@ -93,17 +99,28 @@ async def wait_until_stopped(self) -> None:
await self._proc.wait()

async def run(self) -> None:
await super().run()
await self.start()
self._stopped_future: asyncio.Future = asyncio.Future()

self._proc = proc = await asyncio.create_subprocess_exec(
str(self._driver_executable),
"run-driver",
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=_get_stderr_fileno(),
limit=32768,
)
try:
self._proc = proc = await asyncio.create_subprocess_exec(
str(self._driver_executable),
"run-driver",
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=_get_stderr_fileno(),
limit=32768,
)
except FileNotFoundError:
self.on_error_future.set_exception(
Error(
"playwright's driver is not found, You can read the contributing guide "
"for some guidance on how to get everything setup for working on the code "
"https://github.com/microsoft/playwright-python/blob/master/CONTRIBUTING.md"
)
)
return

assert proc.stdout
assert proc.stdin
self._output = proc.stdin
Expand Down Expand Up @@ -160,15 +177,22 @@ async def wait_until_stopped(self) -> None:
await self._connection.wait_closed()

async def run(self) -> None:
await super().run()
await self.start()

options: Dict[str, Any] = {}
if self.timeout is not None:
options["close_timeout"] = self.timeout / 1000
options["ping_timeout"] = self.timeout / 1000

if self.headers is not None:
options["extra_headers"] = self.headers
self._connection = await websockets.connect(self.ws_endpoint, **options)
try:
self._connection = await websockets.connect(self.ws_endpoint, **options)
except Exception as err:
self.on_error_future.set_exception(
Error(f"playwright's websocket endpoint connection error: {err}")
)
return

while not self._stopped:
try:
Expand Down
16 changes: 14 additions & 2 deletions playwright/async_api/_context_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,22 @@ async def __aenter__(self) -> AsyncPlaywright:
)
loop = asyncio.get_running_loop()
self._connection._loop = loop
obj = asyncio.create_task(
self._connection.wait_for_object_with_known_name("Playwright")
)
await self._connection._transport.start()
loop.create_task(self._connection.run())
playwright = AsyncPlaywright(
await self._connection.wait_for_object_with_known_name("Playwright")
done, pending = await asyncio.wait(
{
obj,
self._connection._transport.on_error_future, # type: ignore
},
return_when=asyncio.FIRST_COMPLETED,
)
if not obj.done():
obj.cancel()
obj = next(iter(done)).result()
playwright = AsyncPlaywright(obj) # type: ignore
playwright.stop = self.__aexit__ # type: ignore
return playwright

Expand Down
10 changes: 10 additions & 0 deletions tests/async/test_browsertype_connect.py
Original file line number Diff line number Diff line change
Expand Up @@ -182,3 +182,13 @@ async def test_prevent_getting_video_path(
== "Path is not available when using browserType.connect(). Use save_as() to save a local copy."
)
remote_server.kill()


async def test_connect_to_closed_server_without_hangs(
browser_type: BrowserType, launch_server
):
remote_server = launch_server()
remote_server.kill()
with pytest.raises(Error) as exc:
await browser_type.connect(remote_server.ws_endpoint)
assert "playwright's websocket endpoint connection error" in exc.value.message
10 changes: 10 additions & 0 deletions tests/sync/test_browsertype_connect.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,3 +145,13 @@ def test_browser_type_connect_should_forward_close_events_to_pages(
assert events == ["page::close", "context::close", "browser::disconnected"]
remote.kill()
assert events == ["page::close", "context::close", "browser::disconnected"]


def test_connect_to_closed_server_without_hangs(
browser_type: BrowserType, launch_server
):
remote_server = launch_server()
remote_server.kill()
with pytest.raises(Error) as exc:
browser_type.connect(remote_server.ws_endpoint)
assert "playwright's websocket endpoint connection error" in exc.value.message