code: def test_default_train_test_simple(
      data_generator: Callable, expected_output: dict
  ):
      """Check if simple splitting perform correctly for all used in FEDOT data types"""
      input_data = data_generator()
      train_data, test_data = train_test_data_setup(
          input_data, stratify=False
      )
  
      assert (
          train_data.features.shape
          == expected_output["train_features_size"]
      )
      assert (
          test_data.features.shape
          == expected_output["test_features_size"]
      )
      assert tuple(test_data.idx) == expected_output["test_idx"]
  
imports: from functools import partial
  from typing import Callable
  import numpy
  import pandas
  import pytest
  from fedot.core.data.data import InputData
  from fedot.core.data.data_split import train_test_data_setup
  from fedot.core.data.multi_modal import MultiModalData
  from fedot.core.optimisers.objective.data_source_splitter import DataSourceSplitter
  from fedot.core.repository.dataset_types import DataTypesEnum
  from fedot.core.repository.tasks import Task, TaskTypesEnum, TsForecastingParams
  from fedot.core.utils import split_data
  from fedot.core.data.cv_folds import cv_generator
  from test.unit.pipelines.test_decompose_pipelines import get_classification_data
  from test.unit.tasks.test_forecasting import get_ts_data
label_desc: Use image node_id 11 for calling a global function with example usage: test_default_train_test_simple(data_generator, expected_output) without return types