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
|
from typing import Optional, Union
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from websockets.exceptions import ConnectionClosedError
|
||||||
|
|
||||||
from .api.api_tokens import list_api_tokens_for_user
|
from .api.api_tokens import list_api_tokens_for_user
|
||||||
from .api.file import (
|
from .api.file import (
|
||||||
@ -381,148 +382,165 @@ def test_ws_simple():
|
|||||||
|
|
||||||
|
|
||||||
def test_ws_import():
|
def test_ws_import():
|
||||||
# Create our client.
|
max_retries = 3
|
||||||
client = ClientFromEnv()
|
for attempt in range(1, max_retries + 1):
|
||||||
|
try:
|
||||||
|
# Create our client.
|
||||||
|
client = ClientFromEnv()
|
||||||
|
|
||||||
# Connect to the websocket.
|
# Connect to the websocket.
|
||||||
with modeling_commands_ws.WebSocket(
|
with modeling_commands_ws.WebSocket(
|
||||||
client=client,
|
client=client,
|
||||||
fps=30,
|
fps=30,
|
||||||
post_effect=PostEffectType.NOEFFECT,
|
post_effect=PostEffectType.NOEFFECT,
|
||||||
show_grid=False,
|
show_grid=False,
|
||||||
unlocked_framerate=False,
|
unlocked_framerate=False,
|
||||||
video_res_height=360,
|
video_res_height=360,
|
||||||
video_res_width=480,
|
video_res_width=480,
|
||||||
webrtc=False,
|
webrtc=False,
|
||||||
) as websocket:
|
) as websocket:
|
||||||
# read the content of the file
|
# read the content of the file
|
||||||
dir_path = os.path.dirname(os.path.realpath(__file__))
|
dir_path = os.path.dirname(os.path.realpath(__file__))
|
||||||
file_name = "ORIGINALVOXEL-3.obj"
|
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)
|
||||||
content = file.read()
|
with open(file_path, "rb") as file:
|
||||||
file.close()
|
content = file.read()
|
||||||
cmd_id = uuid.uuid4()
|
cmd_id = uuid.uuid4()
|
||||||
ImportFile(data=content, path=file_name)
|
ImportFile(data=content, path=file_name)
|
||||||
# form the request
|
# form the request
|
||||||
req = WebSocketRequest(
|
req = WebSocketRequest(
|
||||||
OptionModelingCmdReq(
|
OptionModelingCmdReq(
|
||||||
cmd=ModelingCmd(
|
cmd=ModelingCmd(
|
||||||
OptionImportFiles(
|
OptionImportFiles(
|
||||||
files=[ImportFile(data=content, path=file_name)],
|
files=[ImportFile(data=content, path=file_name)],
|
||||||
format=InputFormat(
|
format=InputFormat(
|
||||||
OptionObj(
|
OptionObj(
|
||||||
units=UnitLength.M,
|
units=UnitLength.M,
|
||||||
coords=System(
|
coords=System(
|
||||||
forward=AxisDirectionPair(
|
forward=AxisDirectionPair(
|
||||||
axis=Axis.Y, direction=Direction.NEGATIVE
|
axis=Axis.Y,
|
||||||
),
|
direction=Direction.NEGATIVE,
|
||||||
up=AxisDirectionPair(
|
),
|
||||||
axis=Axis.Z, direction=Direction.POSITIVE
|
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.
|
# Get the success message.
|
||||||
cmd_id = uuid.uuid4()
|
for message in websocket:
|
||||||
# form the request
|
message_dict = message.model_dump()
|
||||||
req = WebSocketRequest(
|
if message_dict["success"] is not True:
|
||||||
OptionModelingCmdReq(
|
raise Exception(message_dict)
|
||||||
cmd=ModelingCmd(OptionDefaultCameraFocusOn(uuid=object_id)),
|
elif message_dict["resp"]["type"] != "modeling":
|
||||||
cmd_id=ModelingCmdId(cmd_id),
|
continue
|
||||||
)
|
elif (
|
||||||
)
|
message_dict["resp"]["data"]["modeling_response"]["type"]
|
||||||
websocket.send(req)
|
!= "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.
|
# Now we want to focus on the object.
|
||||||
for message in websocket:
|
cmd_id = uuid.uuid4()
|
||||||
message_dict = message.model_dump()
|
# form the request
|
||||||
if message_dict["success"] is not True:
|
req = WebSocketRequest(
|
||||||
raise Exception(message_dict)
|
OptionModelingCmdReq(
|
||||||
elif message_dict["resp"]["type"] != "modeling":
|
cmd=ModelingCmd(OptionDefaultCameraFocusOn(uuid=object_id)),
|
||||||
continue
|
cmd_id=ModelingCmdId(cmd_id),
|
||||||
elif message_dict["request_id"] == str(cmd_id):
|
)
|
||||||
# We got a success response for our cmd.
|
)
|
||||||
break
|
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:
|
else:
|
||||||
raise Exception(message_dict)
|
# After max retries, re-raise the exception to fail the test
|
||||||
|
print(
|
||||||
# Now we want to snapshot as a png.
|
f"ConnectionClosedError encountered on attempt {attempt}/{max_retries}. No more retries left."
|
||||||
cmd_id = uuid.uuid4()
|
)
|
||||||
# form the request
|
raise
|
||||||
# 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)
|
|
||||||
|
|
||||||
|
|
||||||
def test_serialize_deserialize():
|
def test_serialize_deserialize():
|
||||||
|
|||||||
Reference in New Issue
Block a user