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:
Greg Sweeney
GitHub
@ -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,148 +382,165 @@ def test_ws_simple():
|
||||
|
||||
|
||||
def test_ws_import():
|
||||
# Create our client.
|
||||
client = ClientFromEnv()
|
||||
max_retries = 3
|
||||
for attempt in range(1, max_retries + 1):
|
||||
try:
|
||||
# Create our client.
|
||||
client = ClientFromEnv()
|
||||
|
||||
# Connect to the websocket.
|
||||
with modeling_commands_ws.WebSocket(
|
||||
client=client,
|
||||
fps=30,
|
||||
post_effect=PostEffectType.NOEFFECT,
|
||||
show_grid=False,
|
||||
unlocked_framerate=False,
|
||||
video_res_height=360,
|
||||
video_res_width=480,
|
||||
webrtc=False,
|
||||
) as websocket:
|
||||
# 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")
|
||||
content = file.read()
|
||||
file.close()
|
||||
cmd_id = uuid.uuid4()
|
||||
ImportFile(data=content, path=file_name)
|
||||
# form the request
|
||||
req = WebSocketRequest(
|
||||
OptionModelingCmdReq(
|
||||
cmd=ModelingCmd(
|
||||
OptionImportFiles(
|
||||
files=[ImportFile(data=content, path=file_name)],
|
||||
format=InputFormat(
|
||||
OptionObj(
|
||||
units=UnitLength.M,
|
||||
coords=System(
|
||||
forward=AxisDirectionPair(
|
||||
axis=Axis.Y, direction=Direction.NEGATIVE
|
||||
),
|
||||
up=AxisDirectionPair(
|
||||
axis=Axis.Z, direction=Direction.POSITIVE
|
||||
),
|
||||
# Connect to the websocket.
|
||||
with modeling_commands_ws.WebSocket(
|
||||
client=client,
|
||||
fps=30,
|
||||
post_effect=PostEffectType.NOEFFECT,
|
||||
show_grid=False,
|
||||
unlocked_framerate=False,
|
||||
video_res_height=360,
|
||||
video_res_width=480,
|
||||
webrtc=False,
|
||||
) as websocket:
|
||||
# read the content of the file
|
||||
dir_path = os.path.dirname(os.path.realpath(__file__))
|
||||
file_name = "ORIGINALVOXEL-3.obj"
|
||||
file_path = os.path.join(dir_path, "..", "assets", file_name)
|
||||
with open(file_path, "rb") as file:
|
||||
content = file.read()
|
||||
cmd_id = uuid.uuid4()
|
||||
ImportFile(data=content, path=file_name)
|
||||
# form the request
|
||||
req = WebSocketRequest(
|
||||
OptionModelingCmdReq(
|
||||
cmd=ModelingCmd(
|
||||
OptionImportFiles(
|
||||
files=[ImportFile(data=content, path=file_name)],
|
||||
format=InputFormat(
|
||||
OptionObj(
|
||||
units=UnitLength.M,
|
||||
coords=System(
|
||||
forward=AxisDirectionPair(
|
||||
axis=Axis.Y,
|
||||
direction=Direction.NEGATIVE,
|
||||
),
|
||||
up=AxisDirectionPair(
|
||||
axis=Axis.Z,
|
||||
direction=Direction.POSITIVE,
|
||||
),
|
||||
),
|
||||
)
|
||||
),
|
||||
)
|
||||
),
|
||||
cmd_id=ModelingCmdId(cmd_id),
|
||||
)
|
||||
),
|
||||
cmd_id=ModelingCmdId(cmd_id),
|
||||
)
|
||||
)
|
||||
# Import files request must be sent as binary, because the file contents might be binary.
|
||||
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:
|
||||
raise Exception(message_dict)
|
||||
elif message_dict["resp"]["type"] != "modeling":
|
||||
continue
|
||||
elif (
|
||||
message_dict["resp"]["data"]["modeling_response"]["type"]
|
||||
!= "import_files"
|
||||
):
|
||||
# We have a modeling command response.
|
||||
# Make sure its the import files response.
|
||||
raise Exception(message_dict)
|
||||
else:
|
||||
# Okay we have the import files response.
|
||||
# Break since now we know it was a success.
|
||||
object_id = str(
|
||||
message_dict["resp"]["data"]["modeling_response"]["data"][
|
||||
"object_id"
|
||||
]
|
||||
)
|
||||
break
|
||||
# Import files request must be sent as binary, because the file contents might be binary.
|
||||
websocket.send_binary(req)
|
||||
|
||||
# Now we want to focus on the object.
|
||||
cmd_id = uuid.uuid4()
|
||||
# form the request
|
||||
req = WebSocketRequest(
|
||||
OptionModelingCmdReq(
|
||||
cmd=ModelingCmd(OptionDefaultCameraFocusOn(uuid=object_id)),
|
||||
cmd_id=ModelingCmdId(cmd_id),
|
||||
)
|
||||
)
|
||||
websocket.send(req)
|
||||
# Get the success message.
|
||||
for message in websocket:
|
||||
message_dict = message.model_dump()
|
||||
if message_dict["success"] is not True:
|
||||
raise Exception(message_dict)
|
||||
elif message_dict["resp"]["type"] != "modeling":
|
||||
continue
|
||||
elif (
|
||||
message_dict["resp"]["data"]["modeling_response"]["type"]
|
||||
!= "import_files"
|
||||
):
|
||||
# We have a modeling command response.
|
||||
# Make sure its the import files response.
|
||||
raise Exception(message_dict)
|
||||
else:
|
||||
# Okay we have the import files response.
|
||||
# Break since now we know it was a success.
|
||||
object_id = str(
|
||||
message_dict["resp"]["data"]["modeling_response"]["data"][
|
||||
"object_id"
|
||||
]
|
||||
)
|
||||
break
|
||||
|
||||
# Get the success message.
|
||||
for message in websocket:
|
||||
message_dict = message.model_dump()
|
||||
if message_dict["success"] is not True:
|
||||
raise Exception(message_dict)
|
||||
elif message_dict["resp"]["type"] != "modeling":
|
||||
continue
|
||||
elif message_dict["request_id"] == str(cmd_id):
|
||||
# We got a success response for our cmd.
|
||||
break
|
||||
# Now we want to focus on the object.
|
||||
cmd_id = uuid.uuid4()
|
||||
# form the request
|
||||
req = WebSocketRequest(
|
||||
OptionModelingCmdReq(
|
||||
cmd=ModelingCmd(OptionDefaultCameraFocusOn(uuid=object_id)),
|
||||
cmd_id=ModelingCmdId(cmd_id),
|
||||
)
|
||||
)
|
||||
websocket.send(req)
|
||||
|
||||
# Get the success message.
|
||||
for message in websocket:
|
||||
message_dict = message.model_dump()
|
||||
if message_dict["success"] is not True:
|
||||
raise Exception(message_dict)
|
||||
elif message_dict["resp"]["type"] != "modeling":
|
||||
continue
|
||||
elif message_dict["request_id"] == str(cmd_id):
|
||||
# We got a success response for our cmd.
|
||||
break
|
||||
else:
|
||||
raise Exception(message_dict)
|
||||
|
||||
# Now we want to snapshot as a png.
|
||||
cmd_id = uuid.uuid4()
|
||||
# form the request
|
||||
req = WebSocketRequest(
|
||||
OptionModelingCmdReq(
|
||||
cmd=ModelingCmd(OptionTakeSnapshot(format=ImageFormat.PNG)),
|
||||
cmd_id=ModelingCmdId(cmd_id),
|
||||
)
|
||||
)
|
||||
websocket.send(req)
|
||||
|
||||
# Get the success message.
|
||||
for message in websocket:
|
||||
message_dict = message.model_dump()
|
||||
if message_dict["success"] is not True:
|
||||
raise Exception(message_dict)
|
||||
elif message_dict["resp"]["type"] != "modeling":
|
||||
continue
|
||||
elif (
|
||||
message_dict["resp"]["data"]["modeling_response"]["type"]
|
||||
!= "take_snapshot"
|
||||
):
|
||||
# Make sure its the correct response.
|
||||
raise Exception(message_dict)
|
||||
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"]
|
||||
break
|
||||
|
||||
# Save the contents to a file.
|
||||
png_path = os.path.join(dir_path, "..", "assets", "snapshot.png")
|
||||
with open(png_path, "wb") as f:
|
||||
f.write(png_contents)
|
||||
|
||||
# Ensure the file is not empty.
|
||||
assert len(png_contents) > 0
|
||||
|
||||
# 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:
|
||||
raise Exception(message_dict)
|
||||
|
||||
# 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)),
|
||||
cmd_id=ModelingCmdId(cmd_id),
|
||||
)
|
||||
)
|
||||
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:
|
||||
raise Exception(message_dict)
|
||||
elif message_dict["resp"]["type"] != "modeling":
|
||||
continue
|
||||
elif (
|
||||
message_dict["resp"]["data"]["modeling_response"]["type"]
|
||||
!= "take_snapshot"
|
||||
):
|
||||
# Make sure its the correct response.
|
||||
raise Exception(message_dict)
|
||||
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"]
|
||||
break
|
||||
|
||||
# Save the contents to a file.
|
||||
png_path = os.path.join(dir_path, "..", "assets", "snapshot.png")
|
||||
with open(png_path, "wb") as f:
|
||||
f.write(png_contents)
|
||||
|
||||
# Ensure the file is not empty.
|
||||
assert len(png_contents) > 0
|
||||
|
||||
# Ensure the file exists.
|
||||
assert os.path.exists(png_path)
|
||||
# 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():
|
||||
|
||||
Reference in New Issue
Block a user