import unittest

import numpy as np

from v2.aug import get_transforms


class AugTest(unittest.TestCase):
    @staticmethod
    def make_images():
        img = (np.random.rand(100, 100, 3) * 255).astype('uint8')
        return img.copy(), img.copy()

    def test_aug(self):
        for scope in ('strong', 'weak'):
            for crop in ('random', 'center'):
                aug_pipeline = get_transforms(80, scope=scope, crop=crop)
                a, b = self.make_images()
                a, b = aug_pipeline(a, b)
                np.testing.assert_allclose(a, b)