code: def test_gen_edge_files(
      num_chunks, num_parts, edges_fmt, edges_delimiter
  ):
      """Unit test case for the function
      tools/distpartitioning/parmetis_preprocess.py::gen_edge_files
  
      Parameters:
      -----------
      num_chunks : int
          no. of chunks the input graph needs to be partititioned into
      num_parts : int
          no. of partitions
      edges_fmt : string
          specifying the storage format for the edge files
      edges_delimiter : string
          specifying the delimiter used in the edge files
      """
      # Create the input dataset
      with tempfile.TemporaryDirectory() as root_dir:
          # Create expected environment for test
          input_dir = os.path.join(root_dir, "chunked-data")
          output_dir = os.path.join(root_dir, "preproc_dir")
  
          # Mock a parser object
          fn_params = namedtuple(
              "fn_params", "input_dir output_dir num_parts"
          )
          fn_params.input_dir = input_dir
          fn_params.output_dir = output_dir
          fn_params.num_parts = num_parts
  
          # Create test files and get corresponding file schema
          schema_map = _get_test_data(
              input_dir, num_chunks, edges_fmt, edges_delimiter
          )
          edges_file_list = schema_map["edges"][EDGE_TYPE]["data"]
          # This is breaking encapsulation, but no other good way to get file list
          rank_assignments = generate_roundrobin_read_list(
              len(edges_file_list), num_parts
          )
  
          # Get the global node id offsets for each node type
          # There is only one node-type in the test graph
          # which range from 0 thru 9.
          ntype_gnid_offset = {}
          ntype_gnid_offset[NODE_TYPE] = np.array(
              [0, 10 * num_chunks]
          ).reshape(1, 2)
  
          # Iterate over no. of partitions
          for rank in range(num_parts):
              actual_results = gen_edge_files(
                  rank, schema_map, fn_params
              )
  
              # Get the original files
              original_files = [
                  edges_file_list[file_idx]
                  for file_idx in rank_assignments[rank]
              ]
  
              # Validate the results with the baseline results
              # Test 1. no. of files should have the same count per rank
              assert len(original_files) == len(actual_results)
              assert len(actual_results) > 0
  
              # Test 2. Check the contents of each file and verify the
              # file contents match with the expected results.
              for actual_fname, original_fname in zip(
                  actual_results, original_files
              ):
                  # Check the actual file exists
                  assert os.path.isfile(actual_fname)
                  # Read both files and compare the edges
                  # Here note that the src and dst end points are global_node_ids
                  actual_data = _read_file(actual_fname, "csv", " ")
                  expected_data = _read_file(
                      original_fname, edges_fmt, edges_delimiter
                  )
  
                  # Subtract the global node id offsets, so that we get type node ids
                  # In the current unit test case, the graph has only one node-type.
                  # and this means that type-node-ids are same as the global-node-ids.
                  # Below two lines will take take into effect when the graphs have
                  # more than one node type.
                  actual_data[:, 0] -= ntype_gnid_offset[NODE_TYPE][
                      0, 0
                  ]
                  actual_data[:, 1] -= ntype_gnid_offset[NODE_TYPE][
                      0, 0
                  ]
  
                  # Verify that the contents are equal
                  assert_array_equal(expected_data, actual_data)
  
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 3 for calling a global function with example usage: test_gen_edge_files(num_chunks, num_parts, edges_fmt, edges_delimiter) without return types