import os

import unittest

from models.unet import get_unet_model

from train import load_data, load_dataset, shuffling

class TestClassifier(unittest.TestCase):

    def test_model(self):

        """

        Check the input layer and the output layer of the model.

        Args:

            None

        Returns:

            None

        """

        model = get_unet_model((512, 512, 3))

        # Check the image height

        self.assertEqual(

            model.get_layer("input_1").input_shape[0][0],

            tuple(model.get_layer("conv2d_18").output.shape)[0],

        )

        # Check the image width

        self.assertEqual(

            model.get_layer("input_1").input_shape[0][1],

            tuple(model.get_layer("conv2d_18").output.shape)[1],

        )

    def test_dataloader(self):

        """

        Check the number of images and masks of the training and validation pipeline.

        Args:

            None

        Returns:

            None

        """

        dataset_path = os.path.join("new_data")

        train_path = os.path.join(dataset_path, "train")

        valid_path = os.path.join(dataset_path, "valid")

        train_x, train_y = load_data(train_path)

        train_x, train_y = shuffling(train_x, train_y)

        valid_x, valid_y = load_data(valid_path)

        train_dataset = load_dataset(train_x, train_y, batch_size=16)

        valid_dataset = load_dataset(valid_x, valid_y, batch_size=16)

        train_image, train_mask = next(iter(train_dataset))

        valid_image, valid_mask = next(iter(valid_dataset))

        self.assertEqual(len(train_image), len(train_mask))

        self.assertEqual(len(valid_image), len(valid_mask))

if __name__ == "__main__":

    unittest.main()