Skip to content

Commit

Permalink
refactor: Improve graph_state_model.json_schema unit test readability…
Browse files Browse the repository at this point in the history
… and structure.
  • Loading branch information
ogabrielluiz committed Aug 12, 2024
1 parent 2a65e17 commit fb82107
Showing 1 changed file with 50 additions and 16 deletions.
66 changes: 50 additions & 16 deletions src/backend/tests/unit/graph/graph/test_graph_state_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,24 +116,58 @@ def test_graph_state_model_json_schema():

graph = Graph(chat_input, chat_output)
graph.prepare()
# Now iterate through the graph
# and check that the graph is running
# correctly

GraphStateModel = create_state_model_from_graph(graph)
graph_state_model: BaseModel = GraphStateModel()
json_schema = graph_state_model.model_json_schema(mode="serialization")

# Test main schema structure
assert json_schema["title"] == "GraphStateModel"
assert json_schema["type"] == "object"

assert "chat_input" in json_schema["properties"]
assert json_schema["properties"]["chat_input"]["title"] == "ChatInputStateModel"
assert json_schema["properties"]["chat_input"]["type"] == "object"

assert "chat_output" in json_schema["properties"]
assert json_schema["properties"]["chat_output"]["title"] == "ChatOutputStateModel"
assert json_schema["properties"]["chat_output"]["type"] == "object"

assert "Message" in json_schema["$defs"]
assert "Image" in json_schema["$defs"]

assert json_schema["required"] == ["chat_input", "chat_output"]
assert set(json_schema["required"]) == {"chat_input", "chat_output"}

# Test chat_input and chat_output properties
for prop in ["chat_input", "chat_output"]:
assert prop in json_schema["properties"]
assert json_schema["properties"][prop]["allOf"][0]["$ref"].startswith("#/$defs/")
assert json_schema["properties"][prop]["readOnly"] is True

# Test $defs
assert set(json_schema["$defs"].keys()) == {"ChatInputStateModel", "ChatOutputStateModel", "Image", "Message"}

# Test ChatInputStateModel and ChatOutputStateModel
for model in ["ChatInputStateModel", "ChatOutputStateModel"]:
assert json_schema["$defs"][model]["type"] == "object"
assert json_schema["$defs"][model]["title"] == model
assert "message" in json_schema["$defs"][model]["properties"]
assert json_schema["$defs"][model]["properties"]["message"]["allOf"][0]["$ref"] == "#/$defs/Message"
assert json_schema["$defs"][model]["properties"]["message"]["readOnly"] is True
assert json_schema["$defs"][model]["required"] == ["message"]

# Test Message model
message_props = json_schema["$defs"]["Message"]["properties"]
assert set(message_props.keys()) == {
"text_key",
"data",
"default_value",
"text",
"sender",
"sender_name",
"files",
"session_id",
"timestamp",
"flow_id",
}
assert message_props["text_key"]["type"] == "string"
assert message_props["data"]["type"] == "object"
assert "anyOf" in message_props["default_value"]
assert "anyOf" in message_props["files"]
assert message_props["timestamp"]["type"] == "string"

# Test Image model
image_props = json_schema["$defs"]["Image"]["properties"]
assert set(image_props.keys()) == {"path", "url"}
for prop in ["path", "url"]:
assert "anyOf" in image_props[prop]
assert {"type": "string"} in image_props[prop]["anyOf"]
assert {"type": "null"} in image_props[prop]["anyOf"]

0 comments on commit fb82107

Please sign in to comment.