Fix tests/end2end-tests/data_process/test_data.py

awslabs · Nov 15, 2023 · 5f7ee71 · 5f7ee71
1 parent 25f4113
commit 5f7ee71
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/end2end-tests/data_process/test_data.py b/tests/end2end-tests/data_process/test_data.py
@@ -58,9 +58,9 @@ def read_data_parquet(data_file):
 else:
     raise ValueError('Invalid graph format: {}'.format(args.graph_format))
 
-node1_map = read_data_parquet(os.path.join(out_dir, "node1_id_remap.parquet"))
+node1_map = read_data_parquet(os.path.join(out_dir, "node_id_mappings", "node1"))
 reverse_node1_map = {val: key for key, val in zip(node1_map['orig'], node1_map['new'])}
-node3_map = read_data_parquet(os.path.join(out_dir, "node3_id_remap.parquet"))
+node3_map = read_data_parquet(os.path.join(out_dir, "node_id_mappings", "node3"))
 reverse_node3_map = {val: key for key, val in zip(node3_map['orig'], node3_map['new'])}
 
 # Test the first node data