Add retry logic to test for file import via websocket (#274)
* retry `test_ws_import()` up to three attempts * format * little bit of cleanup
This commit is contained in:
@ -5,6 +5,7 @@ import uuid
|
||||
from typing import Optional, Union
|
||||
|
||||
import pytest
|
||||
from websockets.exceptions import ConnectionClosedError
|
||||
|
||||
from .api.api_tokens import list_api_tokens_for_user
|
||||
from .api.file import (
|
||||
@ -381,6 +382,9 @@ def test_ws_simple():
|
||||
|
||||
|
||||
def test_ws_import():
|
||||
max_retries = 3
|
||||
for attempt in range(1, max_retries + 1):
|
||||
try:
|
||||
# Create our client.
|
||||
client = ClientFromEnv()
|
||||
|
||||
@ -398,9 +402,9 @@ def test_ws_import():
|
||||
# read the content of the file
|
||||
dir_path = os.path.dirname(os.path.realpath(__file__))
|
||||
file_name = "ORIGINALVOXEL-3.obj"
|
||||
file = open(os.path.join(dir_path, "..", "assets", file_name), "rb")
|
||||
file_path = os.path.join(dir_path, "..", "assets", file_name)
|
||||
with open(file_path, "rb") as file:
|
||||
content = file.read()
|
||||
file.close()
|
||||
cmd_id = uuid.uuid4()
|
||||
ImportFile(data=content, path=file_name)
|
||||
# form the request
|
||||
@ -414,10 +418,12 @@ def test_ws_import():
|
||||
units=UnitLength.M,
|
||||
coords=System(
|
||||
forward=AxisDirectionPair(
|
||||
axis=Axis.Y, direction=Direction.NEGATIVE
|
||||
axis=Axis.Y,
|
||||
direction=Direction.NEGATIVE,
|
||||
),
|
||||
up=AxisDirectionPair(
|
||||
axis=Axis.Z, direction=Direction.POSITIVE
|
||||
axis=Axis.Z,
|
||||
direction=Direction.POSITIVE,
|
||||
),
|
||||
),
|
||||
)
|
||||
@ -431,7 +437,6 @@ def test_ws_import():
|
||||
websocket.send_binary(req)
|
||||
|
||||
# Get the success message.
|
||||
object_id = ""
|
||||
for message in websocket:
|
||||
message_dict = message.model_dump()
|
||||
if message_dict["success"] is not True:
|
||||
@ -482,7 +487,6 @@ def test_ws_import():
|
||||
# Now we want to snapshot as a png.
|
||||
cmd_id = uuid.uuid4()
|
||||
# form the request
|
||||
# form the request
|
||||
req = WebSocketRequest(
|
||||
OptionModelingCmdReq(
|
||||
cmd=ModelingCmd(OptionTakeSnapshot(format=ImageFormat.PNG)),
|
||||
@ -492,7 +496,6 @@ def test_ws_import():
|
||||
websocket.send(req)
|
||||
|
||||
# Get the success message.
|
||||
png_contents = b""
|
||||
for message in websocket:
|
||||
message_dict = message.model_dump()
|
||||
if message_dict["success"] is not True:
|
||||
@ -508,9 +511,9 @@ def test_ws_import():
|
||||
else:
|
||||
# Okay we have the snapshot response.
|
||||
# Break since now we know it was a success.
|
||||
png_contents = message_dict["resp"]["data"]["modeling_response"][
|
||||
"data"
|
||||
]["contents"]
|
||||
png_contents = message_dict["resp"]["data"][
|
||||
"modeling_response"
|
||||
]["data"]["contents"]
|
||||
break
|
||||
|
||||
# Save the contents to a file.
|
||||
@ -524,6 +527,21 @@ def test_ws_import():
|
||||
# Ensure the file exists.
|
||||
assert os.path.exists(png_path)
|
||||
|
||||
# Exit the retry loop on success
|
||||
break
|
||||
|
||||
except ConnectionClosedError:
|
||||
if attempt < max_retries:
|
||||
print(
|
||||
f"ConnectionClosedError encountered on attempt {attempt}/{max_retries}. Retrying..."
|
||||
)
|
||||
else:
|
||||
# After max retries, re-raise the exception to fail the test
|
||||
print(
|
||||
f"ConnectionClosedError encountered on attempt {attempt}/{max_retries}. No more retries left."
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
def test_serialize_deserialize():
|
||||
json_str = """{"success":true,"request_id":"16a06065-6ca3-4a96-a042-d0bec6b161a6","resp":{"type":"modeling","data":{"modeling_response":{"type":"import_files","data":{"object_id":"f61ac02e-77bd-468f-858f-fd4141a26acd"}}}}}"""
|
||||
|
Reference in New Issue
Block a user