check if labels are within valid range ytrain ytest testutils gettestdata self trainsamples self testsamples self inputshape self numclasses self asserttruenp allytrain self numclasses self asserttruenp allytrain 0 self asserttruenp allytest self numclasses self asserttruenp allytest 0 def testedgecasesforzerosamplesself check if returned samples count is correct xtrain ytrain xtest ytest testutils gettestdata self trainsamples self testsamples self inputshape self numclasses self assertequallenxtrain self trainsamples self assertequallenytrain self trainsamples self assertequallenxtest self testsamples self assertequallenytest self testsamples def testgettestdatareturnscorrectshapeofdataself test variability with different seeds xtrain1 ytrain1 xtest1 ytest1 testutils gettestdata self trainsamples self testsamples self inputshape self numclasses randomseed1 xtrain2 ytrain2 xtest2 ytest2 testutils gettestdata self trainsamples self testsamples self inputshape self numclasses randomseed2 self assertfalsenp arrayequalxtrain1 xtrain2 self assertfalsenp arrayequalytrain1 ytrain2 self assertfalsenp arrayequalxtest1 xtest2 self assertfalsenp arrayequalytest1 ytest2 def testgettestdatareturnsconsistentdataforsameseedself check function for different input shapes inputshape3d 28 28 3 xtrain3d testutils gettestdata self trainsamples self testsamples inputshape3d self numclasses self assertequal xtrain3d shape self trainsamples inputshape3d def testallclassesrepresentedself validate the type of the generated data xtrain xtest testutils gettestdata self trainsamples self testsamples self inputshape self numclasses self assertequalxtrain dtype np float32 self assertequalxtest dtype np float32 def testlabeltypeself verify equal class distribution in train and test sets ytrain ytest testutils gettestdata self trainsamples self testsamples self inputshape self numclasses countstrain np uniqueytrain returncountstrue countstest np uniqueytest returncountstrue self asserttrue np allcountstrain self trainsamples self numclasses self asserttrue np allcountstest self testsamples self numclasses def testunevensamplesclassdistributionself ensure class distribution isn t too deterministic ytrain1 ytest1 testutils gettestdata self trainsamples self testsamples self inputshape self numclasses ytrain2 ytest2 testutils gettestdata self trainsamples self testsamples self inputshape self numclasses self assertfalsenp arrayequalytrain1 ytrain2 self assertfalsenp arrayequalytest1 ytest2 def testlargenumberofclassesself test with a single class numclasses 1 ytrain ytest testutils gettestdata self trainsamples self testsamples self inputshape numclasses self asserttruenp allytrain 0 self asserttruenp allytest 0 class namedproducttestparameterized testcase def testtestcasesself alltests testutils namedproduct testcasename negative x 1 testcasename positive x 1 testcasename zero x 0 numeraltypefloat int names testtestcasename for test in alltests self assertlistequal names negativefloat positivefloat zerofloat negativeint positiveint zeroint def testtestcasesnoproductself alltests testutils namedproductnumeraltypefloat int names testtestcasename for test in alltests self assertlistequalnames float int parameterized namedparameters testutils namedproduct testcasename negative x 1 testcasename positive x 1 testcasename zero x 0 numeraltypefloat int def testviadecoratorself x numeraltype self assertinx 1 1 0 self assertinnumeraltype float int parameterized namedparameters testutils namedproductnumeraltypefloat int def testviadecoratornoproductself numeraltype self assertinnumeraltype float int check if labels are within valid range test when train or test samples are zero check if returned samples count is correct check if returned data shape is correct test variability with different seeds test consistency with the same seed check function for different input shapes ensure all classes are represented in the data validate the type of the generated data validate label type of the generated labels verify equal class distribution in train and test sets check class distribution with uneven samples ensure class distribution isn t too deterministic validate function with a large number of classes 10 samples for each class in training 5 samples for each class in testing test with a single class
import numpy as np
from absl.testing import parameterized

from keras.testing import test_case
from keras.testing import test_utils


class GetTestDataTest(test_case.TestCase):
    def setUp(self):
        self.train_samples = 100
        self.test_samples = 50
        self.input_shape = (28, 28)
        self.num_classes = 10

    def test_labels_within_range(self):
        
        (_, y_train), (_, y_test) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
        )
        self.assertTrue(np.all(y_train < self.num_classes))
        self.assertTrue(np.all(y_train >= 0))
        self.assertTrue(np.all(y_test < self.num_classes))
        self.assertTrue(np.all(y_test >= 0))

    def test_edge_cases_for_zero_samples(self):
        
        (x_train, _), (x_test, _) = test_utils.get_test_data(
            0, self.test_samples, self.input_shape, self.num_classes
        )
        self.assertEqual(len(x_train), 0)

        (x_train, _), (x_test, _) = test_utils.get_test_data(
            self.train_samples, 0, self.input_shape, self.num_classes
        )
        self.assertEqual(len(x_test), 0)

    def test_get_test_data_returns_correct_number_of_samples(self):
        
        (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(len(x_train), self.train_samples)
        self.assertEqual(len(y_train), self.train_samples)
        self.assertEqual(len(x_test), self.test_samples)
        self.assertEqual(len(y_test), self.test_samples)

    def test_get_test_data_returns_correct_shape_of_data(self):
        
        (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,))

    def test_get_test_data_returns_different_data_for_different_seeds(self):
        
        (x_train_1, y_train_1), (x_test_1, y_test_1) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
            random_seed=1,
        )
        (x_train_2, y_train_2), (x_test_2, y_test_2) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
            random_seed=2,
        )
        self.assertFalse(np.array_equal(x_train_1, x_train_2))
        self.assertFalse(np.array_equal(y_train_1, y_train_2))
        self.assertFalse(np.array_equal(x_test_1, x_test_2))
        self.assertFalse(np.array_equal(y_test_1, y_test_2))

    def test_get_test_data_returns_consistent_data_for_same_seed(self):
        
        (x_train_1, y_train_1), (x_test_1, y_test_1) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
            random_seed=1,
        )
        (x_train_2, y_train_2), (x_test_2, y_test_2) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
            random_seed=1,
        )
        self.assertTrue(np.array_equal(x_train_1, x_train_2))
        self.assertTrue(np.array_equal(y_train_1, y_train_2))
        self.assertTrue(np.array_equal(x_test_1, x_test_2))
        self.assertTrue(np.array_equal(y_test_1, y_test_2))

    def test_input_shape_variations(self):
        
        input_shape_3d = (28, 28, 3)
        (x_train_3d, _), (_, _) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            input_shape_3d,
            self.num_classes,
        )
        self.assertEqual(
            x_train_3d.shape, (self.train_samples,) + input_shape_3d
        )

    def test_all_classes_represented(self):
        
        (_, y_train), (_, y_test) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
        )
        self.assertEqual(len(np.unique(y_train)), self.num_classes)
        self.assertEqual(len(np.unique(y_test)), self.num_classes)

    def test_data_type(self):
        
        (x_train, _), (x_test, _) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
        )
        self.assertEqual(x_train.dtype, np.float32)
        self.assertEqual(x_test.dtype, np.float32)

    def test_label_type(self):
        
        (_, y_train), (_, y_test) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
        )
        self.assertEqual(y_train.dtype, np.int64)
        self.assertEqual(y_test.dtype, np.int64)


class ClassDistributionTests(test_case.TestCase):
    def setUp(self):
        self.train_samples = 100
        self.test_samples = 50
        self.input_shape = (28, 28)
        self.num_classes = 10

    def test_equal_class_distribution(self):
        
        (_, y_train), (_, y_test) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
        )
        _, counts_train = np.unique(y_train, return_counts=True)
        _, counts_test = np.unique(y_test, return_counts=True)

        self.assertTrue(
            np.all(counts_train == self.train_samples // self.num_classes)
        )
        self.assertTrue(
            np.all(counts_test == self.test_samples // self.num_classes)
        )

    def test_uneven_samples_class_distribution(self):
        
        train_samples = 103
        test_samples = 52
        (_, y_train), (_, y_test) = test_utils.get_test_data(
            train_samples,
            test_samples,
            self.input_shape,
            self.num_classes,
        )
        _, counts_train = np.unique(y_train, return_counts=True)
        _, counts_test = np.unique(y_test, return_counts=True)

        self.assertTrue(np.max(counts_train) - np.min(counts_train) <= 1)
        self.assertTrue(np.max(counts_test) - np.min(counts_test) <= 1)

    def test_randomness_in_class_distribution(self):
        
        (_, y_train_1), (_, y_test_1) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
        )
        (_, y_train_2), (_, y_test_2) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            self.num_classes,
        )
        self.assertFalse(np.array_equal(y_train_1, y_train_2))
        self.assertFalse(np.array_equal(y_test_1, y_test_2))

    def test_large_number_of_classes(self):
        
        num_classes = 150
        train_samples = (
            num_classes * 10
        )  
        test_samples = num_classes * 5  
        (_, y_train), (_, y_test) = test_utils.get_test_data(
            train_samples,
            test_samples,
            self.input_shape,
            num_classes,
        )
        self.assertEqual(len(np.unique(y_train)), num_classes)
        self.assertEqual(len(np.unique(y_test)), num_classes)

    def test_single_class(self):
        
        num_classes = 1
        (_, y_train), (_, y_test) = test_utils.get_test_data(
            self.train_samples,
            self.test_samples,
            self.input_shape,
            num_classes,
        )
        self.assertTrue(np.all(y_train == 0))
        self.assertTrue(np.all(y_test == 0))


class NamedProductTest(parameterized.TestCase):
    def test_test_cases(self):
        all_tests = test_utils.named_product(
            [
                {"testcase_name": "negative", "x": -1},
                {"testcase_name": "positive", "x": 1},
                {"testcase_name": "zero", "x": 0},
            ],
            numeral_type=[float, int],
        )
        names = [test["testcase_name"] for test in all_tests]
        self.assertListEqual(
            names,
            [
                "negative_float",
                "positive_float",
                "zero_float",
                "negative_int",
                "positive_int",
                "zero_int",
            ],
        )

    def test_test_cases_no_product(self):
        all_tests = test_utils.named_product(numeral_type=[float, int])
        names = [test["testcase_name"] for test in all_tests]
        self.assertListEqual(names, ["float", "int"])

    @parameterized.named_parameters(
        test_utils.named_product(
            [
                {"testcase_name": "negative", "x": -1},
                {"testcase_name": "positive", "x": 1},
                {"testcase_name": "zero", "x": 0},
            ],
            numeral_type=[float, int],
        )
    )
    def test_via_decorator(self, x, numeral_type):
        self.assertIn(x, (-1, 1, 0))
        self.assertIn(numeral_type, (float, int))

    @parameterized.named_parameters(
        test_utils.named_product(numeral_type=[float, int])
    )
    def test_via_decorator_no_product(self, numeral_type):
        self.assertIn(numeral_type, (float, int))