From 8eae2fa1261efb79be6da5505dc0a34c385b1f6f Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 9 Oct 2023 10:41:54 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=8D=20test(graph.py):=20add=20assertio?= =?UTF-8?q?ns=20to=20verify=20expected=20keywords=20in=20source=20and=20ta?= =?UTF-8?q?rget=20fields=20of=20processed=5Fflow=20edges?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The test_process_flow_vector_store_grouped function now includes assertions to check if the expected keywords are present in the source and target fields of the edges in the processed_flow dictionary. This ensures that the processing of the flow vector store grouped JSON is correctly generating the edges with the expected connections between nodes. --- tests/test_graph.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_graph.py b/tests/test_graph.py index 8d442ce76..1a24e0e4b 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -419,6 +419,22 @@ def test_process_flow_vector_store_grouped(vector_store_grouped_json_flow): assert isinstance(processed_flow, dict) assert "nodes" in processed_flow assert "edges" in processed_flow + edges = processed_flow["edges"] + # Expected keywords in source and target fields + expected_keywords = [ + {"source": "VectorStoreInfo", "target": "VectorStoreAgent"}, + {"source": "ChatOpenAI", "target": "VectorStoreAgent"}, + {"source": "OpenAIEmbeddings", "target": "Chroma"}, + {"source": "Chroma", "target": "VectorStoreInfo"}, + {"source": "WebBaseLoader", "target": "RecursiveCharacterTextSplitter"}, + {"source": "RecursiveCharacterTextSplitter", "target": "Chroma"}, + ] + + for idx, expected_keyword in enumerate(expected_keywords): + for key, value in expected_keyword.items(): + assert ( + value in edges[idx][key].split("-")[0] + ), f"Edge {idx}, key {key} expected to contain {value} but got {edges[idx][key]}" def test_update_template(sample_template, sample_nodes):