code: def _test_split_data_by_edge_type(self, g, es_obj, method):
      p = 0.1
      (
          g_test,
          edge_data_ids_test,
          edge_data_labels_test,
      ) = es_obj.train_test_split(
          p=p, method=method, edge_label="friend", keep_connected=True
      )
  
      # if all goes well, what are the expected return values?
      num_sampled_positives = np.sum(edge_data_labels_test == 1)
      num_sampled_negatives = np.sum(edge_data_labels_test == 0)
  
      assert len(edge_data_ids_test) == len(edge_data_labels_test)
      assert (num_sampled_positives - num_sampled_negatives) == 0
      assert len(g_test.edges()) < len(g.edges())
      assert nx.is_connected(g_test)
  
      with pytest.raises(Exception):
          # This call will raise an exception because the graph has no edges of type 'Non Label'
          es_obj.train_test_split(
              p=p,
              method=method,
              keep_connected=True,
              edge_label="No Label",
          )
  
      p = 0.8
      with pytest.raises(ValueError):
          es_obj.train_test_split(
              p=p,
              method=method,
              edge_label="friend",
              keep_connected=True,
          )
  
      p = 0.8
      with pytest.raises(ValueError):
          es_obj.train_test_split(
              p=p,
              method=method,
              edge_label="friend",
              keep_connected=False,
          )
  
imports: import pytest
  import os
  import numpy
  import networkx
  from stellargraph import StellarGraph
  from stellargraph.data.edge_splitter import EdgeSplitter
  from stellargraph.data.epgm import EPGM
  import random
  import datetime
  from datetime import datetime, timedelta
  from ..test_utils import flaky_xfail_mark
  from ..test_utils.graphs import example_graph_random
label_desc: Use image node_id 4 for calling the TestEdgeSplitterHeterogeneous obj's underlying member method code with example usage: obj._test_split_data_by_edge_type(g, es_obj, method) without return types