text: import unittest
  from mock import Mock, patch
  
  import torchbearer.variational.datasets as ds
  
  
  class TestMakeDataset(unittest.TestCase):
      @patch('os.walk')
      def test_make_dataset_all_files(self, mock_walk):
          mock_walk.return_value = [
              ('/root', (), ('b.a', 'c.a'))
          ]
  
          out = ds.make_dataset('a',['a'])
          self.assertTrue(out == ['/root/b.a', '/root/c.a'])
  
      @patch('os.walk')
      def test_make_dataset_some_files(self, mock_walk):
          mock_walk.return_value = [
              ('/root', (), ('b.v', 'c.a'))
          ]
  
          out = ds.make_dataset('a',['a'])
          self.assertTrue(out == ['/root/c.a'])
  
      @patch('os.walk')
      def test_make_dataset_no_files(self, mock_walk):
          mock_walk.return_value = [
              ('/root', (), ('b.v', 'c.v'))
          ]
  
          out = ds.make_dataset('a',['a'])
          self.assertTrue(out == [])
  
      @patch('os.walk')
      def test_make_dataset_no_root(self, mock_walk):
          mock_walk.return_value = [
              ('', (), ('b.a', 'c.a'))
          ]
  
          out = ds.make_dataset('a',['a'])
          self.assertTrue(out == ['b.a', 'c.a'])
  
      @patch('os.walk')
      def test_make_dataset_multiple_extensions(self, mock_walk):
          mock_walk.return_value = [
              ('', (), ('b.1', 'c.2'))
          ]
  
          out = ds.make_dataset('a',['1','2'])
          self.assertTrue(out == ['b.1', 'c.2'])
  
  
  class TestSimpleImageFolder(unittest.TestCase):
      @patch.object(ds, 'make_dataset')
      def test_length(self, mock_md):
          loader = Mock()
          extensions = None
          transform = Mock()
          target_transform = Mock()
  
          images = ['a.1', 'b.1', 'c.1', 'd.1']
          mock_md.return_value = images
          dataset = ds.SimpleImageFolder('root', loader, extensions, transform, target_transform)
  
          self.assertTrue(len(dataset) == len(images))
  
      @patch.object(ds, 'make_dataset')
      def test_image_path(self, mock_md):
          loader = Mock()
          extensions = None
          transform = Mock()
          target_transform = Mock()
  
          images = ['a.1', 'b.1', 'c.1', 'd.1']
          mock_md.return_value = images
          dataset = ds.SimpleImageFolder('root', loader, extensions, transform, target_transform)
  
          _ = dataset[1]
          self.assertTrue(loader.call_args[0][0] == images[1])
  
          _ = dataset[2]
          self.assertTrue(loader.call_args[0][0] == images[2])
  
      @patch.object(ds, 'make_dataset')
      def test_transform_args(self, mock_md):
          loader = Mock()
          mock_image = Mock()
          loader.return_value = mock_image
          extensions = None
          transform = Mock()
          target_transform = Mock()
  
          images = ['a.1', 'b.1', 'c.1', 'd.1']
          mock_md.return_value = images
          dataset = ds.SimpleImageFolder('root', loader, extensions, transform, target_transform)
  
          _ = dataset[1]
          self.assertTrue(transform.call_args[0][0] == mock_image)
  
      @patch.object(ds, 'make_dataset')
      def test_target_transform_args(self, mock_md):
          loader = Mock()
          mock_image = Mock()
          loader.return_value = mock_image
          extensions = None
          transform = Mock()
          target_transform = Mock()
  
          images = ['a.1', 'b.1', 'c.1', 'd.1']
          mock_md.return_value = images
          dataset = ds.SimpleImageFolder('root', loader, extensions, transform, target_transform)
  
          _ = dataset[1]
          self.assertTrue(target_transform.call_args[0][0] == mock_image)
  
  
  class TestCelebA(unittest.TestCase):
      @patch.object(ds.SimpleImageFolder, '__len__')
      @patch.object(ds.SimpleImageFolder, '__getitem__')
      def test_celeba(self, mock_gi, mock_len):
          return_val = 'test'
          mock_len.return_value = 10
          mock_gi.return_value = return_val
  
          celeba = ds.CelebA('root', Mock(), Mock())
          out = celeba[6]
          self.assertTrue(mock_gi.call_args[0][0] == 6)
          self.assertTrue(out == return_val)
  
      @patch.object(ds.SimpleImageFolder, '__len__')
      @patch.object(ds.SimpleImageFolder, '__getitem__')
      def test_celeba_hq_get_item(self, mock_gi, mock_len):
          return_val = 'test'
          mock_len.return_value = 10
          mock_gi.return_value = return_val
  
          celebahq = ds.CelebA_HQ('root', False, Mock())
          out = celebahq[6]
  
          self.assertTrue(mock_gi.call_args[0][0] == 6)
          self.assertTrue(out == return_val)
  
      @patch.object(ds.CelebA_HQ, 'npy_loader')
      @patch.object(ds.SimpleImageFolder, '__len__')
      @patch.object(ds.SimpleImageFolder, '__getitem__')
      def test_celeba_hq_use_numpy_loader(self, mock_gi, mock_len, mock_loader):
          mock_len.return_value = 10
          mock_gi.return_value = Mock()
  
          celebahq = ds.CelebA_HQ('root', True, Mock())
          out = celebahq[6]
  
          self.assertTrue(celebahq.loader == mock_loader)
          self.assertTrue(celebahq.extensions == ['npy'])
  
      @patch.object(ds.SimpleImageFolder, '__len__')
      @patch.object(ds.SimpleImageFolder, '__getitem__')
      @patch('torchbearer.variational.datasets.default_loader')
      def test_celeba_hq_use_default_loader(self, mock_loader, mock_gi, mock_len, ):
          import torchvision
          mock_len.return_value = 10
  
          celebahq = ds.CelebA_HQ('root', False, Mock())
          out = celebahq[6]
  
          self.assertTrue(celebahq.loader == mock_loader)
          self.assertTrue(celebahq.extensions == torchvision.datasets.folder.IMG_EXTENSIONS)
  
      @patch.object(ds.SimpleImageFolder, '__len__')
      @patch('numpy.load')
      @patch('PIL.Image.fromarray')
      def test_celeba_hq_numpy_loader(self, mock_fromarray, mock_load, mock_len):
          samples = ['a.npy', 'b.npy', 'c.npy']
          mock_len.return_value = 3
          mock_output = Mock()
          mock_output.transpose.return_value = mock_output
          mock_load.return_value = (mock_output, )
  
          celebahq = ds.CelebA_HQ('root', True, Mock())
          celebahq.samples = samples
          out = celebahq[1]
  
          self.assertTrue(mock_load.call_args[0][0] == samples[1])
          self.assertTrue(mock_fromarray.call_args[0][0] == mock_output)
  
  
  class TestdSprites(unittest.TestCase):
      @patch.object(ds.dSprites, 'download')
      @patch.object(ds.dSprites, 'load_data')
      @patch('numpy.load')
      def test_no_download(self, mock_load, mock_load_data, mock_download):
          sprites = ds.dSprites('root', False, None)
          self.assertTrue(sprites.download.call_count == 0)
  
      @patch.object(ds.dSprites, 'download')
      @patch.object(ds.dSprites, 'load_data')
      @patch('numpy.load')
      def test_do_download(self, mock_load, mock_load_data, mock_download):
          sprites = ds.dSprites('root', True, None)
          self.assertTrue(sprites.download.call_count == 1)
  
      @patch.object(ds.dSprites, 'download')
      @patch.object(ds.dSprites, 'load_data')
      @patch('numpy.load')
      def test_latents_bases(self, mock_load, mock_load_data, mock_download):
          sprites = ds.dSprites('root', False, None)
          self.assertTrue(list(sprites.latents_bases) == [737280, 245760, 40960, 1024, 32, 1])
  
      @patch.object(ds.dSprites, 'download')
      @patch('numpy.load')
      def test_load_data(self, mock_load, mock_download):
          sprites = ds.dSprites('root', False, None)
          self.assertTrue(mock_load.call_args_list[0][0][0] == 'root/imgs.npy')
          self.assertTrue(mock_load.call_args_list[1][0][0] == 'root/latents_values.npy')
          self.assertTrue(mock_load.call_args_list[2][0][0] == 'root/latents_classes.npy')
  
      @patch.object(ds.dSprites, 'download')
      @patch.object(ds.dSprites, '__getitem__')
      @patch('numpy.load')
      def test_image_by_latent(self, mock_load, mock_gi, mock_download):
          import numpy as np
          mock_gi.return_value = (None, )
  
          bases = np.array([737280, 245760, 40960, 1024, 32, 1])
          latent = [1,2,3,4,5,6]
          img_id = (bases * latent).sum()
  
          sprites = ds.dSprites('root', False, None)
          sprites.get_img_by_latent(latent)
          self.assertTrue(sprites.__getitem__.call_args[0][0] == img_id)
  
      @patch.object(ds.dSprites, 'download')
      @patch.object(ds.dSprites, 'load_data')
      @patch('numpy.load')
      def test_len(self, mock_load, mock_load_data, mock_download):
          import numpy as np
          mock_load_data.return_value = np.zeros((10, 5, 5))
          sprites = ds.dSprites('root', False, None)
  
          self.assertTrue(len(sprites) == 10)
  
      @patch('PIL.Image.fromarray')
      @patch.object(ds.dSprites, 'load_data')
      @patch('numpy.load')
      def test_get_item(self, mock_load, mock_load_data, mock_fromarray):
          import torch
          sprites = ds.dSprites('root', False, None)
          mock_data = torch.rand(5, 1)
          mock_data[1] *= 3
          sprites.data = mock_data
          out = sprites[1]
  
          self.assertTrue(mock_fromarray.call_args[0][0] == mock_data[1]*255)
          self.assertTrue(mock_fromarray.call_args[1] == {'mode':'L'})
          self.assertTrue(len(out) == 2)
  
      @patch('PIL.Image.fromarray')
      @patch.object(ds.dSprites, 'load_data')
      @patch('numpy.load')
      def test_get_item_transform(self, mock_load, mock_load_data, mock_fromarray):
          import torch
          transform = Mock()
          sprites = ds.dSprites('root', False, transform)
          mock_fromarray.return_value = 'test'
          mock_data = torch.rand(5, 1)
          mock_data[1] *= 3
          sprites.data = mock_data
          _ = sprites[1]
  
          self.assertTrue(transform.call_args[0][0] == 'test')
  
      @patch('zipfile.ZipFile')
      @patch('shutil.copyfileobj')
      @patch('os.makedirs')
      @patch.object(ds.dSprites, 'load_data')
      @patch('numpy.load')
      def test_download(self, mock_load, mock_load_data, mock_mkdirs, mock_copyfileobj, mock_zip, mock_open, mock_urlopen):
          mock_return = Mock()
          # mock_urlopen = Mock()
          # mock_url.return_value = mock_urlopen
          mock_urlopen.return_value.__enter__.return_value = mock_return
  
          root = 'root'
          filename = "dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz"
  
          sprites = ds.dSprites(root, True, None)
  
          self.assertTrue(mock_mkdirs.call_args[0][0] == "root")
          self.assertTrue(mock_mkdirs.call_args[1] == {'exist_ok': True})
  
          self.assertTrue(mock_urlopen.call_args[0][0] == 'https://github.com/deepmind/dsprites-dataset/blob/master/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz?raw=true')
          self.assertTrue(mock_open.call_args[0][0] == 'root/' + filename)
  
          self.assertTrue(mock_copyfileobj.call_args[0][0] == mock_return)
  
      import sys
      if sys.version_info[0] < 3:
          test_download = patch('urllib2.urlopen')(patch('__builtin__.open')(test_download))
      else:
          test_download = patch('urllib.request.urlopen')(patch('builtins.open')(test_download))