Skip to content

fix(transport): handle connection error correctly (#650) #651

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
May 4, 2021
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