code: def _get_test_data(edges_dir, num_chunks, edge_fmt, edge_fmt_del):
      """Creates unit test input which are a set of edge files
      in the following format "src_node_id<delimiter>dst_node_id"
  
      Parameters:
      -----------
      edges_dir : str
          folder where edge files are stored
      num_chunks : int
          no. of files to create for each edge type
      edge_fmt : str, optional
          to specify whether this file is csv or parquet
      edge_fmt_del : str optional
          delimiter to use in the edges file
  
      Returns:
      --------
      dict :
          dictionary created which represents the schema used for
          creating the input dataset
      """
      schema = {}
      schema["num_nodes_per_type"] = [10]
      schema["edge_type"] = [EDGE_TYPE]
      schema["node_type"] = [NODE_TYPE]
  
      edges = {}
      edges[EDGE_TYPE] = {}
      edges[EDGE_TYPE]["format"] = {}
      edges[EDGE_TYPE]["format"]["name"] = edge_fmt
      edges[EDGE_TYPE]["format"]["delimiter"] = edge_fmt_del
  
      os.makedirs(edges_dir, exist_ok=True)
      fmt_meta = {"name": edge_fmt}
      if edge_fmt == "csv":
          fmt_meta["delimiter"] = edge_fmt_del
  
      edge_files = []
      for idx in range(num_chunks):
          path = os.path.join(
              edges_dir, f"test_file_{idx}.{fmt_meta['name']}"
          )
          array_parser = array_readwriter.get_array_parser(**fmt_meta)
          edge_data = (
              np.array([np.arange(10), np.arange(10)]).reshape(10, 2)
              + 10 * idx
          )
          array_parser.write(path, edge_data)
  
          edge_files.append(path)
  
      edges[EDGE_TYPE]["data"] = edge_files
      schema["edges"] = edges
  
      return schema
  
file_path: dgl/tests/tools/test_parmetis_preproc.py
imports: import os
  import tempfile
  from collections import namedtuple
  import numpy
  import pytest
  from distpartitioning import array_readwriter, constants
  from distpartitioning.parmetis_preprocess import gen_edge_files
  from distpartitioning.utils import generate_roundrobin_read_list
  from numpy.testing import assert_array_equal
label_desc: Use image node_id 2 for calling a global function with example usage: _get_test_data(edges_dir, num_chunks, edge_fmt, edge_fmt_del) and returns: schema