PyORAM/src/pyoram/tests/test_block_storage.py

   1 import os
   2 import shutil
   3 import unittest2
   4 import tempfile
   5 import struct
   6
   7 from pyoram.storage.block_storage import \
   8     BlockStorageTypeFactory
   9 from pyoram.storage.block_storage_file import \
  10      BlockStorageFile
  11 from pyoram.storage.block_storage_mmap import \
  12      BlockStorageMMap
  13 from pyoram.storage.block_storage_ram import \
  14      BlockStorageRAM
  15 from pyoram.storage.block_storage_sftp import \
  16      BlockStorageSFTP
  17 from pyoram.storage.block_storage_s3 import \
  18      BlockStorageS3
  19 from pyoram.storage.boto3_s3_wrapper import \
  20     (Boto3S3Wrapper,
  21      MockBoto3S3Wrapper)
  22
  23 import six
  24 from six.moves import xrange
  25 from six import BytesIO
  26
  27 thisdir = os.path.dirname(os.path.abspath(__file__))
  28
  29 try:
  30     import boto3
  31     has_boto3 = True
  32 except:                                                # pragma: no cover
  33     has_boto3 = False                                  # pragma: no cover
  34
  35 class TestBlockStorageTypeFactory(unittest2.TestCase):
  36
  37     def test_file(self):
  38         self.assertIs(BlockStorageTypeFactory('file'),
  39                       BlockStorageFile)
  40
  41     def test_mmap(self):
  42         self.assertIs(BlockStorageTypeFactory('mmap'),
  43                       BlockStorageMMap)
  44
  45     def test_ram(self):
  46         self.assertIs(BlockStorageTypeFactory('ram'),
  47                       BlockStorageRAM)
  48
  49     def test_sftp(self):
  50         self.assertIs(BlockStorageTypeFactory('sftp'),
  51                       BlockStorageSFTP)
  52
  53     def test_s3(self):
  54         self.assertIs(BlockStorageTypeFactory('s3'),
  55                       BlockStorageS3)
  56
  57     def test_invalid(self):
  58         with self.assertRaises(ValueError):
  59             BlockStorageTypeFactory(None)
  60
  61     def test_register_invalid_name(self):
  62         with self.assertRaises(ValueError):
  63             BlockStorageTypeFactory.register_device(
  64                 's3', BlockStorageFile)
  65
  66     def test_register_invalid_type(self):
  67         with self.assertRaises(TypeError):
  68             BlockStorageTypeFactory.register_device(
  69                 'new_str_type', str)
  70
  71 class _TestBlockStorage(object):
  72
  73     _type = None
  74     _type_kwds = None
  75
  76     @classmethod
  77     def _read_storage(cls, storage):
  78         with open(storage.storage_name, 'rb') as f:
  79             return f.read()
  80
  81     @classmethod
  82     def _remove_storage(cls, name):
  83         if os.path.exists(name):
  84             if os.path.isdir(name):
  85                 shutil.rmtree(name, ignore_errors=True)
  86             else:
  87                 os.remove(name)
  88
  89     @classmethod
  90     def _check_exists(cls, name):
  91         return os.path.exists(name)
  92
  93     @classmethod
  94     def _get_empty_existing(cls):
  95         return os.path.join(thisdir,
  96                             "baselines",
  97                             "exists.empty")
  98
  99     @classmethod
 100     def _get_dummy_noexist(cls):
 101         fd, name = tempfile.mkstemp(dir=os.getcwd())
 102         os.close(fd)
 103         return name
 104
 105     def _open_teststorage(self, **kwds):
 106         kwds.update(self._type_kwds)
 107         return self._type(self._testfname, **kwds)
 108
 109     def _reopen_storage(self, storage):
 110         return self._type(storage.storage_name, **self._type_kwds)
 111
 112     @classmethod
 113     def setUpClass(cls):
 114         assert cls._type is not None
 115         assert cls._type_kwds is not None
 116         cls._dummy_name = cls._get_dummy_noexist()
 117         if cls._check_exists(cls._dummy_name):
 118             cls._remove_storage(cls._dummy_name)
 119         if os.path.exists(cls._dummy_name):
 120             _TestBlockStorage.\
 121                 _remove_storage(cls._dummy_name)       # pragma: no cover
 122         cls._block_size = 25
 123         cls._block_count = 5
 124         cls._testfname = cls.__name__ + "_testfile.bin"
 125         cls._blocks = []
 126         f = cls._type.setup(
 127             cls._testfname,
 128             block_size=cls._block_size,
 129             block_count=cls._block_count,
 130             initialize=lambda i: bytes(bytearray([i])*cls._block_size),
 131             ignore_existing=True,
 132             **cls._type_kwds)
 133         f.close()
 134         cls._original_f = f
 135         for i in range(cls._block_count):
 136             data = bytearray([i])*cls._block_size
 137             cls._blocks.append(data)
 138
 139     @classmethod
 140     def tearDownClass(cls):
 141         cls._remove_storage(cls._testfname)
 142         cls._remove_storage(cls._dummy_name)
 143
 144     def test_setup_fails(self):
 145         self.assertEqual(self._check_exists(self._dummy_name), False)
 146         with self.assertRaises(IOError):
 147             self._type.setup(
 148                 self._get_empty_existing(),
 149                 block_size=10,
 150                 block_count=10,
 151                 **self._type_kwds)
 152         self.assertEqual(self._check_exists(self._dummy_name), False)
 153         with self.assertRaises(IOError):
 154             self._type.setup(
 155                 self._get_empty_existing(),
 156                 block_size=10,
 157                 block_count=10,
 158                 ignore_existing=False,
 159                 **self._type_kwds)
 160         self.assertEqual(self._check_exists(self._dummy_name), False)
 161         with self.assertRaises(ValueError):
 162             self._type.setup(self._dummy_name,
 163                              block_size=0,
 164                              block_count=1,
 165                              **self._type_kwds)
 166         self.assertEqual(self._check_exists(self._dummy_name), False)
 167         with self.assertRaises(ValueError):
 168             self._type.setup(self._dummy_name,
 169                              block_size=1,
 170                              block_count=0,
 171                              **self._type_kwds)
 172         self.assertEqual(self._check_exists(self._dummy_name), False)
 173         with self.assertRaises(TypeError):
 174             self._type.setup(self._dummy_name,
 175                              block_size=1,
 176                              block_count=1,
 177                              header_data=2,
 178                              **self._type_kwds)
 179         self.assertEqual(self._check_exists(self._dummy_name), False)
 180         # TODO: The multiprocessing module is bad
 181         #       about handling exceptions raised on the
 182         #       thread's stack.
 183         #with self.assertRaises(ValueError):
 184         #    def _init(i):
 185         #        raise ValueError
 186         #    self._type.setup(self._dummy_name,
 187         #                     block_size=1,
 188         #                     block_count=1,
 189         #                     initialize=_init,
 190         #                     **self._type_kwds)
 191         #self.assertEqual(self._check_exists(self._dummy_name), False)
 192
 193     def test_setup(self):
 194         fname = ".".join(self.id().split(".")[1:])
 195         fname += ".bin"
 196         fname = os.path.join(thisdir, fname)
 197         self._remove_storage(fname)
 198         bsize = 10
 199         bcount = 11
 200         fsetup = self._type.setup(fname, bsize, bcount, **self._type_kwds)
 201         fsetup.close()
 202         flen = len(self._read_storage(fsetup))
 203         self.assertEqual(
 204             flen,
 205             self._type.compute_storage_size(bsize,
 206                                             bcount))
 207         self.assertEqual(
 208             flen >
 209             self._type.compute_storage_size(bsize,
 210                                             bcount,
 211                                             ignore_header=True),
 212             True)
 213         with self._reopen_storage(fsetup) as f:
 214             self.assertEqual(f.header_data, bytes())
 215             self.assertEqual(fsetup.header_data, bytes())
 216             self.assertEqual(f.block_size, bsize)
 217             self.assertEqual(fsetup.block_size, bsize)
 218             self.assertEqual(f.block_count, bcount)
 219             self.assertEqual(fsetup.block_count, bcount)
 220             self.assertEqual(f.storage_name, fsetup.storage_name)
 221             self.assertEqual(fsetup.storage_name, fsetup.storage_name)
 222             if self._type is not BlockStorageRAM:
 223                 self.assertEqual(fsetup.storage_name, fname)
 224             else:
 225                 self.assertEqual(fsetup.storage_name, None)
 226         self._remove_storage(fname)
 227
 228     def test_setup_withdata(self):
 229         fname = ".".join(self.id().split(".")[1:])
 230         fname += ".bin"
 231         fname = os.path.join(thisdir, fname)
 232         self._remove_storage(fname)
 233         bsize = 10
 234         bcount = 11
 235         header_data = bytes(bytearray([0,1,2]))
 236         fsetup = self._type.setup(fname,
 237                                   bsize,
 238                                   bcount,
 239                                   header_data=header_data,
 240                                   **self._type_kwds)
 241         fsetup.close()
 242
 243         flen = len(self._read_storage(fsetup))
 244         self.assertEqual(
 245             flen,
 246             self._type.compute_storage_size(bsize,
 247                                             bcount,
 248                                             header_data=header_data))
 249         self.assertTrue(len(header_data) > 0)
 250         self.assertEqual(
 251             self._type.compute_storage_size(bsize,
 252                                             bcount) <
 253             self._type.compute_storage_size(bsize,
 254                                             bcount,
 255                                             header_data=header_data),
 256             True)
 257         self.assertEqual(
 258             flen >
 259             self._type.compute_storage_size(bsize,
 260                                             bcount,
 261                                             header_data=header_data,
 262                                             ignore_header=True),
 263             True)
 264         with self._reopen_storage(fsetup) as f:
 265             self.assertEqual(f.header_data, header_data)
 266             self.assertEqual(fsetup.header_data, header_data)
 267             self.assertEqual(f.block_size, bsize)
 268             self.assertEqual(fsetup.block_size, bsize)
 269             self.assertEqual(f.block_count, bcount)
 270             self.assertEqual(fsetup.block_count, bcount)
 271             self.assertEqual(f.storage_name, fsetup.storage_name)
 272             self.assertEqual(fsetup.storage_name, fsetup.storage_name)
 273             if self._type is not BlockStorageRAM:
 274                 self.assertEqual(fsetup.storage_name, fname)
 275             else:
 276                 self.assertEqual(fsetup.storage_name, None)
 277
 278         self._remove_storage(fname)
 279
 280     def test_init_noexists(self):
 281         self.assertEqual(self._check_exists(self._dummy_name), False)
 282         with self.assertRaises(IOError):
 283             with self._type(self._dummy_name, **self._type_kwds) as f:
 284                 pass                                   # pragma: no cover
 285
 286     def test_init_exists(self):
 287         self.assertEqual(self._check_exists(self._testfname), True)
 288         databefore = self._read_storage(self._original_f)
 289         with self._open_teststorage() as f:
 290             self.assertEqual(f.block_size, self._block_size)
 291             self.assertEqual(f.block_count, self._block_count)
 292             self.assertEqual(f.storage_name, self._testfname)
 293             self.assertEqual(f.header_data, bytes())
 294         self.assertEqual(self._check_exists(self._testfname), True)
 295         dataafter = self._read_storage(self._original_f)
 296         self.assertEqual(databefore, dataafter)
 297
 298     def test_read_block(self):
 299         with self._open_teststorage() as f:
 300             self.assertEqual(f.bytes_sent, 0)
 301             self.assertEqual(f.bytes_received, 0)
 302
 303             for i, data in enumerate(self._blocks):
 304                 self.assertEqual(list(bytearray(f.read_block(i))),
 305                                  list(self._blocks[i]))
 306             for i, data in enumerate(self._blocks):
 307                 self.assertEqual(list(bytearray(f.read_block(i))),
 308                                  list(self._blocks[i]))
 309             for i, data in reversed(list(enumerate(self._blocks))):
 310                 self.assertEqual(list(bytearray(f.read_block(i))),
 311                                  list(self._blocks[i]))
 312             for i, data in reversed(list(enumerate(self._blocks))):
 313                 self.assertEqual(list(bytearray(f.read_block(i))),
 314                                  list(self._blocks[i]))
 315
 316             self.assertEqual(f.bytes_sent, 0)
 317             self.assertEqual(f.bytes_received,
 318                              self._block_count*self._block_size*4)
 319
 320         with self._open_teststorage() as f:
 321             self.assertEqual(f.bytes_sent, 0)
 322             self.assertEqual(f.bytes_received, 0)
 323
 324             self.assertEqual(list(bytearray(f.read_block(0))),
 325                              list(self._blocks[0]))
 326             self.assertEqual(list(bytearray(f.read_block(self._block_count-1))),
 327                              list(self._blocks[-1]))
 328
 329             self.assertEqual(f.bytes_sent, 0)
 330             self.assertEqual(f.bytes_received,
 331                              self._block_size*2)
 332
 333     def test_write_block(self):
 334         data = bytearray([self._block_count])*self._block_size
 335         self.assertEqual(len(data) > 0, True)
 336         with self._open_teststorage() as f:
 337             self.assertEqual(f.bytes_sent, 0)
 338             self.assertEqual(f.bytes_received, 0)
 339
 340             for i in xrange(self._block_count):
 341                 self.assertNotEqual(list(bytearray(f.read_block(i))),
 342                                     list(data))
 343             for i in xrange(self._block_count):
 344                 f.write_block(i, bytes(data))
 345             for i in xrange(self._block_count):
 346                 self.assertEqual(list(bytearray(f.read_block(i))),
 347                                  list(data))
 348             for i, block in enumerate(self._blocks):
 349                 f.write_block(i, bytes(block))
 350
 351             self.assertEqual(f.bytes_sent,
 352                              self._block_count*self._block_size*2)
 353             self.assertEqual(f.bytes_received,
 354                              self._block_count*self._block_size*2)
 355
 356     def test_read_blocks(self):
 357         with self._open_teststorage() as f:
 358             self.assertEqual(f.bytes_sent, 0)
 359             self.assertEqual(f.bytes_received, 0)
 360
 361             data = f.read_blocks(list(xrange(self._block_count)))
 362             self.assertEqual(len(data), self._block_count)
 363             for i, block in enumerate(data):
 364                 self.assertEqual(list(bytearray(block)),
 365                                  list(self._blocks[i]))
 366             data = f.read_blocks([0])
 367             self.assertEqual(len(data), 1)
 368             self.assertEqual(list(bytearray(data[0])),
 369                              list(self._blocks[0]))
 370             self.assertEqual(len(self._blocks) > 1, True)
 371             data = f.read_blocks(list(xrange(1, self._block_count)) + [0])
 372             self.assertEqual(len(data), self._block_count)
 373             for i, block in enumerate(data[:-1], 1):
 374                 self.assertEqual(list(bytearray(block)),
 375                                  list(self._blocks[i]))
 376             self.assertEqual(list(bytearray(data[-1])),
 377                              list(self._blocks[0]))
 378
 379             self.assertEqual(f.bytes_sent, 0)
 380             self.assertEqual(f.bytes_received,
 381                              (2*self._block_count+1)*self._block_size)
 382
 383     def test_yield_blocks(self):
 384         with self._open_teststorage() as f:
 385             self.assertEqual(f.bytes_sent, 0)
 386             self.assertEqual(f.bytes_received, 0)
 387
 388             data = list(f.yield_blocks(list(xrange(self._block_count))))
 389             self.assertEqual(len(data), self._block_count)
 390             for i, block in enumerate(data):
 391                 self.assertEqual(list(bytearray(block)),
 392                                  list(self._blocks[i]))
 393             data = list(f.yield_blocks([0]))
 394             self.assertEqual(len(data), 1)
 395             self.assertEqual(list(bytearray(data[0])),
 396                              list(self._blocks[0]))
 397             self.assertEqual(len(self._blocks) > 1, True)
 398             data = list(f.yield_blocks(list(xrange(1, self._block_count)) + [0]))
 399             self.assertEqual(len(data), self._block_count)
 400             for i, block in enumerate(data[:-1], 1):
 401                 self.assertEqual(list(bytearray(block)),
 402                                  list(self._blocks[i]))
 403             self.assertEqual(list(bytearray(data[-1])),
 404                              list(self._blocks[0]))
 405
 406             self.assertEqual(f.bytes_sent, 0)
 407             self.assertEqual(f.bytes_received,
 408                              (2*self._block_count+1)*self._block_size)
 409
 410     def test_write_blocks(self):
 411         data = [bytearray([self._block_count])*self._block_size
 412                 for i in xrange(self._block_count)]
 413         with self._open_teststorage() as f:
 414             self.assertEqual(f.bytes_sent, 0)
 415             self.assertEqual(f.bytes_received, 0)
 416
 417             orig = f.read_blocks(list(xrange(self._block_count)))
 418             self.assertEqual(len(orig), self._block_count)
 419             for i, block in enumerate(orig):
 420                 self.assertEqual(list(bytearray(block)),
 421                                  list(self._blocks[i]))
 422             f.write_blocks(list(xrange(self._block_count)),
 423                            [bytes(b) for b in data])
 424             new = f.read_blocks(list(xrange(self._block_count)))
 425             self.assertEqual(len(new), self._block_count)
 426             for i, block in enumerate(new):
 427                 self.assertEqual(list(bytearray(block)),
 428                                  list(data[i]))
 429             f.write_blocks(list(xrange(self._block_count)),
 430                            [bytes(b) for b in self._blocks])
 431             orig = f.read_blocks(list(xrange(self._block_count)))
 432             self.assertEqual(len(orig), self._block_count)
 433             for i, block in enumerate(orig):
 434                 self.assertEqual(list(bytearray(block)),
 435                                  list(self._blocks[i]))
 436
 437             self.assertEqual(f.bytes_sent,
 438                              self._block_count*self._block_size*2)
 439             self.assertEqual(f.bytes_received,
 440                              self._block_count*self._block_size*3)
 441
 442     def test_update_header_data(self):
 443         fname = ".".join(self.id().split(".")[1:])
 444         fname += ".bin"
 445         fname = os.path.join(thisdir, fname)
 446         self._remove_storage(fname)
 447         bsize = 10
 448         bcount = 11
 449         header_data = bytes(bytearray([0,1,2]))
 450         fsetup = self._type.setup(fname,
 451                                   block_size=bsize,
 452                                   block_count=bcount,
 453                                   header_data=header_data,
 454                                   **self._type_kwds)
 455         fsetup.close()
 456         new_header_data = bytes(bytearray([1,1,1]))
 457         with self._reopen_storage(fsetup) as f:
 458             self.assertEqual(f.header_data, header_data)
 459             f.update_header_data(new_header_data)
 460             self.assertEqual(f.header_data, new_header_data)
 461         with self._reopen_storage(fsetup) as f:
 462             self.assertEqual(f.header_data, new_header_data)
 463         with self.assertRaises(ValueError):
 464             with self._reopen_storage(fsetup) as f:
 465                 f.update_header_data(bytes(bytearray([1,1])))
 466         with self.assertRaises(ValueError):
 467             with self._reopen_storage(fsetup) as f:
 468                 f.update_header_data(bytes(bytearray([1,1,1,1])))
 469         with self._reopen_storage(fsetup) as f:
 470             self.assertEqual(f.header_data, new_header_data)
 471         self._remove_storage(fname)
 472
 473     def test_locked_flag(self):
 474         with self._open_teststorage() as f:
 475             with self.assertRaises(IOError):
 476                 with self._open_teststorage() as f1:
 477                     pass                               # pragma: no cover
 478             with self.assertRaises(IOError):
 479                 with self._open_teststorage() as f1:
 480                     pass                               # pragma: no cover
 481             with self._open_teststorage(ignore_lock=True) as f1:
 482                 pass
 483             with self.assertRaises(IOError):
 484                 with self._open_teststorage() as f1:
 485                     pass                               # pragma: no cover
 486             with self._open_teststorage(ignore_lock=True) as f1:
 487                 pass
 488             with self._open_teststorage(ignore_lock=True) as f1:
 489                 pass
 490         with self._open_teststorage(ignore_lock=True) as f:
 491             pass
 492
 493     def test_read_block_cloned(self):
 494         with self._open_teststorage() as forig:
 495             self.assertEqual(forig.bytes_sent, 0)
 496             self.assertEqual(forig.bytes_received, 0)
 497             with forig.clone_device() as f:
 498                 self.assertEqual(forig.bytes_sent, 0)
 499                 self.assertEqual(forig.bytes_received, 0)
 500                 self.assertEqual(f.bytes_sent, 0)
 501                 self.assertEqual(f.bytes_received, 0)
 502
 503                 for i, data in enumerate(self._blocks):
 504                     self.assertEqual(list(bytearray(f.read_block(i))),
 505                                      list(self._blocks[i]))
 506                 for i, data in enumerate(self._blocks):
 507                     self.assertEqual(list(bytearray(f.read_block(i))),
 508                                      list(self._blocks[i]))
 509                 for i, data in reversed(list(enumerate(self._blocks))):
 510                     self.assertEqual(list(bytearray(f.read_block(i))),
 511                                      list(self._blocks[i]))
 512                 for i, data in reversed(list(enumerate(self._blocks))):
 513                     self.assertEqual(list(bytearray(f.read_block(i))),
 514                                      list(self._blocks[i]))
 515
 516                 self.assertEqual(f.bytes_sent, 0)
 517                 self.assertEqual(f.bytes_received,
 518                                  self._block_count*self._block_size*4)
 519
 520             with forig.clone_device() as f:
 521                 self.assertEqual(forig.bytes_sent, 0)
 522                 self.assertEqual(forig.bytes_received, 0)
 523                 self.assertEqual(f.bytes_sent, 0)
 524                 self.assertEqual(f.bytes_received, 0)
 525
 526                 self.assertEqual(list(bytearray(f.read_block(0))),
 527                                  list(self._blocks[0]))
 528                 self.assertEqual(list(bytearray(f.read_block(self._block_count-1))),
 529                                  list(self._blocks[-1]))
 530
 531                 self.assertEqual(f.bytes_sent, 0)
 532                 self.assertEqual(f.bytes_received,
 533                                  self._block_size*2)
 534             self.assertEqual(forig.bytes_sent, 0)
 535             self.assertEqual(forig.bytes_received, 0)
 536
 537     def test_write_block_cloned(self):
 538         data = bytearray([self._block_count])*self._block_size
 539         self.assertEqual(len(data) > 0, True)
 540         with self._open_teststorage() as forig:
 541             self.assertEqual(forig.bytes_sent, 0)
 542             self.assertEqual(forig.bytes_received, 0)
 543             with forig.clone_device() as f:
 544                 self.assertEqual(forig.bytes_sent, 0)
 545                 self.assertEqual(forig.bytes_received, 0)
 546                 self.assertEqual(f.bytes_sent, 0)
 547                 self.assertEqual(f.bytes_received, 0)
 548
 549                 for i in xrange(self._block_count):
 550                     self.assertNotEqual(list(bytearray(f.read_block(i))),
 551                                         list(data))
 552                 for i in xrange(self._block_count):
 553                     f.write_block(i, bytes(data))
 554                 for i in xrange(self._block_count):
 555                     self.assertEqual(list(bytearray(f.read_block(i))),
 556                                      list(data))
 557                 for i, block in enumerate(self._blocks):
 558                     f.write_block(i, bytes(block))
 559
 560                 self.assertEqual(f.bytes_sent,
 561                                  self._block_count*self._block_size*2)
 562                 self.assertEqual(f.bytes_received,
 563                                  self._block_count*self._block_size*2)
 564             self.assertEqual(forig.bytes_sent, 0)
 565             self.assertEqual(forig.bytes_received, 0)
 566
 567     def test_read_blocks_cloned(self):
 568         with self._open_teststorage() as forig:
 569             self.assertEqual(forig.bytes_sent, 0)
 570             self.assertEqual(forig.bytes_received, 0)
 571             with forig.clone_device() as f:
 572                 self.assertEqual(forig.bytes_sent, 0)
 573                 self.assertEqual(forig.bytes_received, 0)
 574                 self.assertEqual(f.bytes_sent, 0)
 575                 self.assertEqual(f.bytes_received, 0)
 576
 577                 data = f.read_blocks(list(xrange(self._block_count)))
 578                 self.assertEqual(len(data), self._block_count)
 579                 for i, block in enumerate(data):
 580                     self.assertEqual(list(bytearray(block)),
 581                                      list(self._blocks[i]))
 582                 data = f.read_blocks([0])
 583                 self.assertEqual(len(data), 1)
 584                 self.assertEqual(list(bytearray(data[0])),
 585                                  list(self._blocks[0]))
 586                 self.assertEqual(len(self._blocks) > 1, True)
 587                 data = f.read_blocks(list(xrange(1, self._block_count)) + [0])
 588                 self.assertEqual(len(data), self._block_count)
 589                 for i, block in enumerate(data[:-1], 1):
 590                     self.assertEqual(list(bytearray(block)),
 591                                      list(self._blocks[i]))
 592                 self.assertEqual(list(bytearray(data[-1])),
 593                                  list(self._blocks[0]))
 594
 595                 self.assertEqual(f.bytes_sent, 0)
 596                 self.assertEqual(f.bytes_received,
 597                                  (2*self._block_count + 1)*self._block_size)
 598             self.assertEqual(forig.bytes_sent, 0)
 599             self.assertEqual(forig.bytes_received, 0)
 600
 601     def test_yield_blocks_cloned(self):
 602         with self._open_teststorage() as forig:
 603             self.assertEqual(forig.bytes_sent, 0)
 604             self.assertEqual(forig.bytes_received, 0)
 605             with forig.clone_device() as f:
 606                 self.assertEqual(forig.bytes_sent, 0)
 607                 self.assertEqual(forig.bytes_received, 0)
 608                 self.assertEqual(f.bytes_sent, 0)
 609                 self.assertEqual(f.bytes_received, 0)
 610
 611                 data = list(f.yield_blocks(list(xrange(self._block_count))))
 612                 self.assertEqual(len(data), self._block_count)
 613                 for i, block in enumerate(data):
 614                     self.assertEqual(list(bytearray(block)),
 615                                      list(self._blocks[i]))
 616                 data = list(f.yield_blocks([0]))
 617                 self.assertEqual(len(data), 1)
 618                 self.assertEqual(list(bytearray(data[0])),
 619                                  list(self._blocks[0]))
 620                 self.assertEqual(len(self._blocks) > 1, True)
 621                 data = list(f.yield_blocks(list(xrange(1, self._block_count)) + [0]))
 622                 self.assertEqual(len(data), self._block_count)
 623                 for i, block in enumerate(data[:-1], 1):
 624                     self.assertEqual(list(bytearray(block)),
 625                                      list(self._blocks[i]))
 626                 self.assertEqual(list(bytearray(data[-1])),
 627                                  list(self._blocks[0]))
 628
 629                 self.assertEqual(f.bytes_sent, 0)
 630                 self.assertEqual(f.bytes_received,
 631                                  (2*self._block_count + 1)*self._block_size)
 632             self.assertEqual(forig.bytes_sent, 0)
 633             self.assertEqual(forig.bytes_received, 0)
 634
 635     def test_write_blocks_cloned(self):
 636         data = [bytearray([self._block_count])*self._block_size
 637                 for i in xrange(self._block_count)]
 638         with self._open_teststorage() as forig:
 639             self.assertEqual(forig.bytes_sent, 0)
 640             self.assertEqual(forig.bytes_received, 0)
 641             with forig.clone_device() as f:
 642                 self.assertEqual(forig.bytes_sent, 0)
 643                 self.assertEqual(forig.bytes_received, 0)
 644                 self.assertEqual(f.bytes_sent, 0)
 645                 self.assertEqual(f.bytes_received, 0)
 646
 647                 orig = f.read_blocks(list(xrange(self._block_count)))
 648                 self.assertEqual(len(orig), self._block_count)
 649                 for i, block in enumerate(orig):
 650                     self.assertEqual(list(bytearray(block)),
 651                                      list(self._blocks[i]))
 652                 f.write_blocks(list(xrange(self._block_count)),
 653                                [bytes(b) for b in data])
 654                 new = f.read_blocks(list(xrange(self._block_count)))
 655                 self.assertEqual(len(new), self._block_count)
 656                 for i, block in enumerate(new):
 657                     self.assertEqual(list(bytearray(block)),
 658                                      list(data[i]))
 659                 f.write_blocks(list(xrange(self._block_count)),
 660                                [bytes(b) for b in self._blocks])
 661                 orig = f.read_blocks(list(xrange(self._block_count)))
 662                 self.assertEqual(len(orig), self._block_count)
 663                 for i, block in enumerate(orig):
 664                     self.assertEqual(list(bytearray(block)),
 665                                      list(self._blocks[i]))
 666
 667                 self.assertEqual(f.bytes_sent,
 668                                  self._block_count*self._block_size*2)
 669                 self.assertEqual(f.bytes_received,
 670                                  self._block_count*self._block_size*3)
 671             self.assertEqual(forig.bytes_sent, 0)
 672             self.assertEqual(forig.bytes_received, 0)
 673
 674 class TestBlockStorageFile(_TestBlockStorage,
 675                            unittest2.TestCase):
 676     _type = BlockStorageFile
 677     _type_kwds = {}
 678
 679 class TestBlockStorageFileNoThreadPool(_TestBlockStorage,
 680                                        unittest2.TestCase):
 681     _type = BlockStorageFile
 682     _type_kwds = {'threadpool_size': 0}
 683
 684 class TestBlockStorageFileThreadPool(_TestBlockStorage,
 685                                      unittest2.TestCase):
 686     _type = BlockStorageFile
 687     _type_kwds = {'threadpool_size': 1}
 688
 689 class TestBlockStorageMMap(_TestBlockStorage,
 690                            unittest2.TestCase):
 691     _type = BlockStorageMMap
 692     _type_kwds = {}
 693
 694 class _TestBlockStorageRAM(_TestBlockStorage):
 695
 696     @classmethod
 697     def _read_storage(cls, storage):
 698         return storage.data
 699
 700     @classmethod
 701     def _remove_storage(cls, name):
 702         pass
 703
 704     @classmethod
 705     def _check_exists(cls, name):
 706         return True
 707
 708     def _open_teststorage(self, **kwds):
 709         kwds.update(self._type_kwds)
 710         return self._type(self._original_f.data, **kwds)
 711
 712     def _reopen_storage(self, storage):
 713         return self._type(storage.data, **self._type_kwds)
 714
 715     #
 716     # Override some of the test methods
 717     #
 718
 719     def test_setup_fails(self):
 720         with self.assertRaises(ValueError):
 721             self._type.setup(self._dummy_name,
 722                              block_size=0,
 723                              block_count=1,
 724                              **self._type_kwds)
 725         with self.assertRaises(ValueError):
 726             self._type.setup(self._dummy_name,
 727                              block_size=1,
 728                              block_count=0,
 729                              **self._type_kwds)
 730         with self.assertRaises(TypeError):
 731             self._type.setup(self._dummy_name,
 732                              block_size=1,
 733                              block_count=1,
 734                              header_data=2,
 735                              **self._type_kwds)
 736
 737     def test_init_noexists(self):
 738         with self.assertRaises(TypeError):
 739             with self._type(2, **self._type_kwds) as f:
 740                 pass                                   # pragma: no cover
 741         with self.assertRaises(TypeError):
 742             with self._type(None, **self._type_kwds) as f:
 743                 pass                                   # pragma: no cover
 744         with self.assertRaises(struct.error):
 745             with self._type(bytearray(), **self._type_kwds) as f:
 746                 pass                                   # pragma: no cover
 747
 748     def test_init_exists(self):
 749         databefore = self._read_storage(self._original_f)
 750         with self._open_teststorage() as f:
 751             self.assertEqual(f.block_size, self._block_size)
 752             self.assertEqual(f.block_count, self._block_count)
 753             self.assertEqual(f.storage_name, self._original_f.storage_name)
 754             self.assertEqual(f.storage_name, None)
 755             self.assertEqual(f.header_data, bytes())
 756         dataafter = self._read_storage(self._original_f)
 757         self.assertEqual(databefore, dataafter)
 758
 759     def test_tofile_fromfile_fileobj(self):
 760         out1 = BytesIO()
 761         self._original_f.tofile(out1)
 762         out1.seek(0)
 763         self.assertEqual(len(self._original_f.data) > 0, True)
 764         self.assertEqual(self._original_f.data, out1.read())
 765         out1.seek(0)
 766         in1 = self._type.fromfile(out1)
 767         self.assertNotEqual(self._original_f.data, in1.data)
 768         out2 = BytesIO()
 769         in1.tofile(out2)
 770         self.assertNotEqual(self._original_f.data, in1.data)
 771         in1.close()
 772         self.assertEqual(self._original_f.data, in1.data)
 773         out2.seek(0)
 774         with self.assertRaises(IOError):
 775             with self._type.fromfile(out2) as in2:
 776                 pass                                  # pragma: no cover
 777         out2.seek(0)
 778         with self._type.fromfile(out2, ignore_lock=True) as in2:
 779             self.assertEqual(self._original_f.data, in1.data)
 780             self.assertNotEqual(self._original_f.data, in2.data)
 781         self.assertEqual(self._original_f.data, in1.data)
 782         self.assertNotEqual(self._original_f.data, in2.data)
 783
 784     def test_tofile_fromfile_filename(self):
 785
 786         def _create():
 787             fd, out = tempfile.mkstemp()
 788             os.close(fd)
 789             return out
 790         def _read(name):
 791             with open(name, 'rb') as f:
 792                 return f.read()
 793
 794         out1 = _create()
 795         self._original_f.tofile(out1)
 796         self.assertEqual(len(self._original_f.data) > 0, True)
 797         self.assertEqual(self._original_f.data, _read(out1))
 798         in1 = self._type.fromfile(out1)
 799         self.assertNotEqual(self._original_f.data, in1.data)
 800         out2 = _create()
 801         in1.tofile(out2)
 802         self.assertNotEqual(self._original_f.data, in1.data)
 803         in1.close()
 804         self.assertEqual(self._original_f.data, in1.data)
 805         with self.assertRaises(IOError):
 806             with self._type.fromfile(out2) as in2:
 807                 pass                                  # pragma: no cover
 808         with self._type.fromfile(out2, ignore_lock=True) as in2:
 809             self.assertEqual(self._original_f.data, in1.data)
 810             self.assertNotEqual(self._original_f.data, in2.data)
 811         self.assertEqual(self._original_f.data, in1.data)
 812         self.assertNotEqual(self._original_f.data, in2.data)
 813
 814 class TestBlockStorageRAM(_TestBlockStorageRAM,
 815                           unittest2.TestCase):
 816     _type = BlockStorageRAM
 817     _type_kwds = {}
 818
 819 class _dummy_sftp_file(object):
 820     def __init__(self, *args, **kwds):
 821         self._f = open(*args, **kwds)
 822     def __enter__(self):
 823         return self
 824     def __exit__(self, *args):
 825         self._f.close()
 826     def readv(self, chunks):
 827         data = []
 828         for offset, size in chunks:
 829             self._f.seek(offset)
 830             data.append(self._f.read(size))
 831         return data
 832     def __getattr__(self, key):
 833         return getattr(self._f, key)
 834     def set_pipelined(self):
 835         pass
 836
 837 class dummy_sftp(object):
 838     remove = os.remove
 839     stat = os.stat
 840     @staticmethod
 841     def open(*args, **kwds):
 842         return _dummy_sftp_file(*args, **kwds)
 843     @staticmethod
 844     def close():
 845         pass
 846
 847 class dummy_sshclient(object):
 848     @staticmethod
 849     def open_sftp():
 850         return dummy_sftp
 851
 852 class TestBlockStorageSFTP(_TestBlockStorage,
 853                            unittest2.TestCase):
 854     _type = BlockStorageSFTP
 855     _type_kwds = {'sshclient': dummy_sshclient}
 856
 857     def test_setup_fails_no_sshclient(self):
 858         self.assertEqual(self._check_exists(self._dummy_name), False)
 859         kwds = dict(self._type_kwds)
 860         del kwds['sshclient']
 861         with self.assertRaises(ValueError):
 862             self._type.setup(self._dummy_name,
 863                              block_size=1,
 864                              block_count=1,
 865                              **kwds)
 866         self.assertEqual(self._check_exists(self._dummy_name), False)
 867
 868     def test_init_exists_no_sshclient(self):
 869         self.assertEqual(self._check_exists(self._testfname), True)
 870         kwds = dict(self._type_kwds)
 871         del kwds['sshclient']
 872         with self.assertRaises(ValueError):
 873             with self._type(self._testfname, **kwds) as f:
 874                 pass                                   # pragma: no cover
 875
 876         databefore = self._read_storage(self._original_f)
 877         with self._open_teststorage() as f:
 878             self.assertEqual(f.block_size, self._block_size)
 879             self.assertEqual(f.block_count, self._block_count)
 880             self.assertEqual(f.storage_name, self._testfname)
 881             self.assertEqual(f.header_data, bytes())
 882         self.assertEqual(self._check_exists(self._testfname), True)
 883         dataafter = self._read_storage(self._original_f)
 884         self.assertEqual(databefore, dataafter)
 885
 886
 887 class _TestBlockStorageS3Mock(_TestBlockStorage):
 888     _type = BlockStorageS3
 889     _type_kwds = {}
 890
 891     @classmethod
 892     def _read_storage(cls, storage):
 893         import glob
 894         data = bytearray()
 895         name = storage.storage_name
 896         prefix_len = len(os.path.join(name,"b"))
 897         nblocks = max(int(bfile[prefix_len:]) for bfile in glob.glob(name+"/b*")) + 1
 898         with open(os.path.join(name, BlockStorageS3._index_name), 'rb') as f:
 899             data.extend(f.read())
 900         for i in range(nblocks):
 901             with open(os.path.join(name, "b"+str(i)), 'rb') as f:
 902                 data.extend(f.read())
 903         return data
 904
 905     def test_init_exists_no_bucket(self):
 906         self.assertEqual(self._check_exists(self._testfname), True)
 907         databefore = self._read_storage(self._original_f)
 908         with self._open_teststorage() as f:
 909             self.assertEqual(f.block_size, self._block_size)
 910             self.assertEqual(f.block_count, self._block_count)
 911             self.assertEqual(f.storage_name, self._testfname)
 912             self.assertEqual(f.header_data, bytes())
 913         self.assertEqual(self._check_exists(self._testfname), True)
 914         dataafter = self._read_storage(self._original_f)
 915         self.assertEqual(databefore, dataafter)
 916         kwds = dict(self._type_kwds)
 917         del kwds['bucket_name']
 918         with self.assertRaises(ValueError):
 919             with self._type(self._testfname, **kwds) as f:
 920                 pass                                   # pragma: no cover
 921         dataafter = self._read_storage(self._original_f)
 922         self.assertEqual(databefore, dataafter)
 923
 924     def test_setup_fails_no_bucket(self):
 925         self.assertEqual(self._check_exists(self._dummy_name), False)
 926         kwds = dict(self._type_kwds)
 927         del kwds['bucket_name']
 928         with self.assertRaises(ValueError):
 929             self._type.setup(self._dummy_name,
 930                              block_size=1,
 931                              block_count=1,
 932                              **kwds)
 933         self.assertEqual(self._check_exists(self._dummy_name), False)
 934
 935     def test_setup_ignore_existing(self):
 936         self.assertEqual(self._check_exists(self._dummy_name), False)
 937         with self._type.setup(self._dummy_name,
 938                               block_size=1,
 939                               block_count=1,
 940                               **self._type_kwds) as f:
 941             pass
 942         self.assertEqual(self._check_exists(self._dummy_name), True)
 943         with self.assertRaises(IOError):
 944             with self._type.setup(self._dummy_name,
 945                                   block_size=1,
 946                                   block_count=1,
 947                                   **self._type_kwds) as f:
 948                 pass                                   # pragma: no cover
 949         self.assertEqual(self._check_exists(self._dummy_name), True)
 950         with self._type.setup(self._dummy_name,
 951                               block_size=1,
 952                               block_count=1,
 953                               ignore_existing=True,
 954                               **self._type_kwds) as f:
 955             pass
 956         self.assertEqual(self._check_exists(self._dummy_name), True)
 957         self._remove_storage(self._dummy_name)
 958
 959 class TestBlockStorageS3Mock(_TestBlockStorageS3Mock,
 960                              unittest2.TestCase):
 961     _type_kwds = {'s3_wrapper': MockBoto3S3Wrapper,
 962                   'bucket_name': '.'}
 963
 964 class TestBlockStorageS3MockNoThreadPool(_TestBlockStorageS3Mock,
 965                              unittest2.TestCase):
 966     _type_kwds = {'s3_wrapper': MockBoto3S3Wrapper,
 967                   'bucket_name': '.',
 968                   'threadpool_size': 0}
 969
 970 class TestBlockStorageS3MockThreadPool(_TestBlockStorageS3Mock,
 971                                        unittest2.TestCase):
 972     _type_kwds = {'s3_wrapper': MockBoto3S3Wrapper,
 973                   'bucket_name': '.',
 974                   'threadpool_size': 4}
 975
 976 @unittest2.skipIf((os.environ.get('PYORAM_AWS_TEST_BUCKET') is None) or \
 977                  (not has_boto3),
 978                  "No PYORAM_AWS_TEST_BUCKET defined in environment or "
 979                  "boto3 is not available")
 980 class TestBlockStorageS3(_TestBlockStorage,
 981                          unittest2.TestCase):
 982     _type = BlockStorageS3
 983     _type_kwds = {'bucket_name': os.environ.get('PYORAM_AWS_TEST_BUCKET')}
 984
 985     @classmethod
 986     def _read_storage(cls, storage):
 987         data = bytearray()
 988         name = storage.storage_name
 989         s3 = Boto3S3Wrapper(cls._type_kwds['bucket_name'])
 990         prefix_len = len(name+"/b")
 991         nblocks = 1 + max(int(obj.key[prefix_len:]) for obj
 992                           in s3._bucket.objects.filter(Prefix=name+"/b"))
 993         data.extend(s3.download(name+"/"+BlockStorageS3._index_name))
 994         for i in range(nblocks):
 995             data.extend(s3.download(name+"/b"+str(i)))
 996         return data
 997
 998     @classmethod
 999     def _remove_storage(cls, name):
1000         Boto3S3Wrapper(cls._type_kwds['bucket_name']).clear(name)
1001
1002     @classmethod
1003     def _check_exists(cls, name):
1004         return Boto3S3Wrapper(cls._type_kwds['bucket_name']).exists(name)
1005
1006     @classmethod
1007     def _get_empty_existing(cls):
1008         return "exists.empty"
1009
1010     @classmethod
1011     def _get_dummy_noexist(cls):
1012         s3 = Boto3S3Wrapper(cls._type_kwds['bucket_name'])
1013         fd, name = tempfile.mkstemp(dir=os.getcwd())
1014         os.close(fd)
1015         os.remove(name)
1016         while s3.exists(name):
1017             fd, name = tempfile.mkstemp(dir=os.getcwd())
1018             os.close(fd)
1019             os.remove(name)
1020         return name
1021
1022 if __name__ == "__main__":
1023     unittest2.main()                                    # pragma: no cover