code: def test_split(ray_start_4_cpus):
      ds = ray.data.range(10)
  
      # Split all by default
      test = TestBasic(
          2,
          True,
          {"train": 5, "test": 5, "val": 5},
          datasets={"train": ds, "test": ds, "val": ds},
      )
      test.fit()
  
      # Test flag "all"
      test = TestBasic(
          2,
          True,
          {"train": 5, "test": 5},
          datasets={"train": ds, "test": ds},
          dataset_config=DataConfig(datasets_to_split="all"),
      )
  
      # Test split train only.
      test = TestBasic(
          2,
          True,
          {"train": 5, "test": 10},
          datasets={"train": ds, "test": ds},
          dataset_config=DataConfig(datasets_to_split=["train"]),
      )
      test.fit()
  
      # Test invalid arguments
      for datasets_to_split in ["train", ("train"), {}]:
          with pytest.raises(
              TypeError, match="`datasets_to_split` should be.*"
          ):
              test = TestBasic(
                  2,
                  True,
                  {"train": 5, "test": 10},
                  datasets={"train": ds, "test": ds},
                  dataset_config=DataConfig(
                      datasets_to_split=datasets_to_split
                  ),
              )
  
      # Test empty `datasets_to_split` list
      test = TestBasic(
          2,
          True,
          {"train": 10, "test": 10},
          datasets={"train": ds, "test": ds},
          dataset_config=DataConfig(datasets_to_split=[]),
      )
      test.fit()
  
imports: from typing import Optional
  import random
  import pytest
  from unittest.mock import MagicMock
  import ray
  from ray import train
  from ray.train import DataConfig, ScalingConfig
  from ray.data import DataIterator
  from ray.train.data_parallel_trainer import DataParallelTrainer
  from ray.data._internal.execution.interfaces.execution_options import ExecutionOptions
  from ray.tests.conftest import *
label_desc: Use image node_id 3 for calling a global function with example usage: test_split(ray_start_4_cpus) without return types