def test_receive(unused_tcp_port, ws):
async def handler(websocket, path):
await websocket.send("Hello World")
ws("localhost", unused_tcp_port, handler)
with ExitStack() as stack:
duplexer = SyncWebsocketDuplexer(
f"ws://localhost:{unused_tcp_port}",
f"ws://localhost:{unused_tcp_port}",
None,
None,
)
stack.callback(duplexer.stop)
assert next(duplexer.receive()) == "Hello World"
def test_unexpected_close(unused_tcp_port, ws):
async def handler(websocket, path):
await websocket.close()
ws("localhost", unused_tcp_port, handler)
with ExitStack() as stack:
duplexer = SyncWebsocketDuplexer(
f"ws://localhost:{unused_tcp_port}",
f"ws://localhost:{unused_tcp_port}",
None,
None,
)
stack.callback(duplexer.stop)
with pytest.raises(ConnectionClosedOK):
next(duplexer.receive())
def test_failed_send(unused_tcp_port, ws):
async def handler(websocket, path):
await websocket.recv()
ws("localhost", unused_tcp_port, handler)
duplexer = SyncWebsocketDuplexer(
f"ws://localhost:{unused_tcp_port}",
f"ws://localhost:{unused_tcp_port}",
None,
None,
)
with patch("websockets.WebSocketCommonProtocol.send") as send:
send.side_effect = OSError("expected OSError")
with pytest.raises(OSError, match="expected OSError"):
duplexer.send("hello")
def track(self):
with ExitStack() as stack:
duplexer = self._ws_duplexer
if not duplexer:
duplexer = SyncWebsocketDuplexer(self._client_uri,
self._base_uri, self._cert,
self._token)
stack.callback(duplexer.stop)
for message in duplexer.receive():
try:
event = from_json(
message,
data_unmarshaller=serialization.evaluator_unmarshaller)
except DataUnmarshallerError:
event = from_json(message, data_unmarshaller=pickle.loads)
yield event
if event["type"] == identifiers.EVTYPE_EE_TERMINATED:
logger.debug(
f"monitor-{self._id} client received terminated")
break
def test_immediate_stop(unused_tcp_port, ws):
# if the duplexer is immediately stopped, it should be well behaved and close
# the connection normally
closed = Event()
async def handler(websocket, path):
try:
await websocket.recv()
except ConnectionClosedOK:
closed.set()
ws("localhost", unused_tcp_port, handler)
duplexer = SyncWebsocketDuplexer(
f"ws://localhost:{unused_tcp_port}",
f"ws://localhost:{unused_tcp_port}",
None,
None,
)
duplexer.stop()
assert closed.wait(10)
def test_secure_echo(unused_tcp_port, ws):
config = EvaluatorServerConfig(unused_tcp_port)
async def handler(websocket, path):
msg = await websocket.recv()
await websocket.send(msg)
ws(
config.host,
unused_tcp_port,
handler,
ssl=config.get_server_ssl_context(),
sock=config.get_socket(),
)
with ExitStack() as stack:
duplexer = SyncWebsocketDuplexer(
f"wss://{config.host}:{unused_tcp_port}",
f"wss://{config.host}:{unused_tcp_port}",
cert=config.cert,
token=None,
)
stack.callback(duplexer.stop)
duplexer.send("Hello Secure World")
assert next(duplexer.receive()) == "Hello Secure World"
def test_generator(unused_tcp_port, ws):
async def handler(websocket, path):
await websocket.send("one")
await websocket.send("two")
await websocket.send("three")
await websocket.send("four")
ws("localhost", unused_tcp_port, handler)
with ExitStack() as stack:
duplexer = SyncWebsocketDuplexer(
f"ws://localhost:{unused_tcp_port}",
f"ws://localhost:{unused_tcp_port}",
None,
None,
)
stack.callback(duplexer.stop)
expected = ["one", "two", "three"]
for msg in duplexer.receive():
assert msg == expected.pop(0)
# Cause a GeneratorExit
if len(expected) == 1:
break
def test_immediate_stop():
duplexer = SyncWebsocketDuplexer("ws://localhost", "", None, None)
duplexer.stop()
def __enter__(self):
self._ws_duplexer = SyncWebsocketDuplexer(self._client_uri,
self._base_uri, self._cert,
self._token)
return self
class _Monitor:
def __init__(self, host, port, protocol="wss", cert=None, token=None):
self._base_uri = f"{protocol}://{host}:{port}"
self._client_uri = f"{self._base_uri}/client"
self._result_uri = f"{self._base_uri}/result"
self._cert = cert
self._token = token
self._ws_duplexer: Optional[SyncWebsocketDuplexer] = None
self._id = str(uuid.uuid1()).split("-")[0]
def __enter__(self):
self._ws_duplexer = SyncWebsocketDuplexer(self._client_uri,
self._base_uri, self._cert,
self._token)
return self
def __exit__(self, *args):
self._ws_duplexer.stop()
def get_base_uri(self):
return self._base_uri
def _send_event(self, cloud_event: CloudEvent) -> None:
with ExitStack() as stack:
duplexer = self._ws_duplexer
if not duplexer:
duplexer = SyncWebsocketDuplexer(self._client_uri,
self._base_uri, self._cert,
self._token)
stack.callback(duplexer.stop)
duplexer.send(
to_json(cloud_event,
data_marshaller=serialization.evaluator_marshaller))
def signal_cancel(self):
logger.debug(f"monitor-{self._id} asking server to cancel...")
out_cloudevent = CloudEvent({
"type": identifiers.EVTYPE_EE_USER_CANCEL,
"source": f"/ert/monitor/{self._id}",
"id": str(uuid.uuid1()),
})
self._send_event(out_cloudevent)
logger.debug(f"monitor-{self._id} asked server to cancel")
def signal_done(self):
logger.debug(f"monitor-{self._id} informing server monitor is done...")
out_cloudevent = CloudEvent({
"type": identifiers.EVTYPE_EE_USER_DONE,
"source": f"/ert/monitor/{self._id}",
"id": str(uuid.uuid1()),
})
self._send_event(out_cloudevent)
logger.debug(f"monitor-{self._id} informed server monitor is done")
def track(self):
with ExitStack() as stack:
duplexer = self._ws_duplexer
if not duplexer:
duplexer = SyncWebsocketDuplexer(self._client_uri,
self._base_uri, self._cert,
self._token)
stack.callback(duplexer.stop)
for message in duplexer.receive():
try:
event = from_json(
message,
data_unmarshaller=serialization.evaluator_unmarshaller)
except DataUnmarshallerError:
event = from_json(message, data_unmarshaller=pickle.loads)
yield event
if event["type"] == identifiers.EVTYPE_EE_TERMINATED:
logger.debug(
f"monitor-{self._id} client received terminated")
break
def test_failed_connection():
with patch("ert.ensemble_evaluator.wait_for_evaluator") as w:
w.side_effect = OSError("expected OSError")
with pytest.raises(OSError, match="expected OSError"):
SyncWebsocketDuplexer("ws://localhost:0", "http://localhost:0",
None, None)
