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
2024-09-16 14:07:27 -04:00
committed by GitHub
parent 5204e30f00
commit c8a3ba33ef

View File

@ -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"}}}}}"""