From c8a3ba33ef0cc6f84ca4fc96c66f6638f9ff1458 Mon Sep 17 00:00:00 2001 From: Greg Sweeney Date: Mon, 16 Sep 2024 14:07:27 -0400 Subject: [PATCH] Add retry logic to test for file import via websocket (#274) * retry `test_ws_import()` up to three attempts * format * little bit of cleanup --- kittycad/client_test.py | 284 +++++++++++++++++++++------------------- 1 file changed, 151 insertions(+), 133 deletions(-) diff --git a/kittycad/client_test.py b/kittycad/client_test.py index 270e32064..c21118238 100644 --- a/kittycad/client_test.py +++ b/kittycad/client_test.py @@ -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():