def test_get_test_data_returns_correct_shape_of_data(self):
        """Check if returned data shape is correct."""
        (x_train, y_train), (x_test, y_test) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
        )
        self.assertEqual(
            x_train.shape, (self.train_samples,) + self.input_shape
        )
        self.assertEqual(y_train.shape, (self.train_samples,))
        self.assertEqual(x_test.shape, (self.test_samples,) + self.input_shape)
        self.assertEqual(y_test.shape, (self.test_samples,))
Check if returned data shape is correct.
keras/src/testing/test_utils_test.py