ByoungSooPark/kittycad.py
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add retry logic to test for file import via websocket (#274)

Browse Source 
* 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
1 changed files with 151 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
kittycad/client_test.py
View File

@ -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,6 +382,9 @@ def test_ws_simple():
def test_ws_import(): def test_ws_import():
max_retries = 3
for attempt in range(1, max_retries + 1):
try:
# Create our client. # Create our client.
client = ClientFromEnv() client = ClientFromEnv()
@ -398,9 +402,9 @@ def test_ws_import():
# 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)
with open(file_path, "rb") as file:
content = file.read() content = file.read()
file.close()
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
@ -414,10 +418,12 @@ def test_ws_import():
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( 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) websocket.send_binary(req)
# Get the success message. # Get the success message.
object_id = ""
for message in websocket: for message in websocket:
message_dict = message.model_dump() message_dict = message.model_dump()
if message_dict["success"] is not True: if message_dict["success"] is not True:
@ -482,7 +487,6 @@ def test_ws_import():
# Now we want to snapshot as a png. # Now we want to snapshot as a png.
cmd_id = uuid.uuid4() cmd_id = uuid.uuid4()
# form the request # form the request
# form the request
req = WebSocketRequest( req = WebSocketRequest(
OptionModelingCmdReq( OptionModelingCmdReq(
cmd=ModelingCmd(OptionTakeSnapshot(format=ImageFormat.PNG)), cmd=ModelingCmd(OptionTakeSnapshot(format=ImageFormat.PNG)),
@ -492,7 +496,6 @@ def test_ws_import():
websocket.send(req) websocket.send(req)
# Get the success message. # Get the success message.
png_contents = b""
for message in websocket: for message in websocket:
message_dict = message.model_dump() message_dict = message.model_dump()
if message_dict["success"] is not True: if message_dict["success"] is not True:
@ -508,9 +511,9 @@ def test_ws_import():
else: else:
# Okay we have the snapshot response. # Okay we have the snapshot response.
# Break since now we know it was a success. # Break since now we know it was a success.
png_contents = message_dict["resp"]["data"]["modeling_response"][ png_contents = message_dict["resp"]["data"][
"data" "modeling_response"
]["contents"] ]["data"]["contents"]
break break
# Save the contents to a file. # Save the contents to a file.
@ -524,6 +527,21 @@ def test_ws_import():
# Ensure the file exists. # Ensure the file exists.
assert os.path.exists(png_path) 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(): 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"}}}}}""" 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"}}}}}"""