7 from pyoram.storage.block_storage import \
8 BlockStorageTypeFactory
9 from pyoram.storage.block_storage_file import \
11 from pyoram.storage.block_storage_mmap import \
13 from pyoram.storage.block_storage_ram import \
15 from pyoram.storage.block_storage_sftp import \
17 from pyoram.storage.block_storage_s3 import \
19 from pyoram.storage.boto3_s3_wrapper import \
24 from six.moves import xrange
25 from six import BytesIO
27 thisdir = os.path.dirname(os.path.abspath(__file__))
32 except: # pragma: no cover
33 has_boto3 = False # pragma: no cover
35 class TestBlockStorageTypeFactory(unittest2.TestCase):
38 self.assertIs(BlockStorageTypeFactory('file'),
42 self.assertIs(BlockStorageTypeFactory('mmap'),
46 self.assertIs(BlockStorageTypeFactory('ram'),
50 self.assertIs(BlockStorageTypeFactory('sftp'),
54 self.assertIs(BlockStorageTypeFactory('s3'),
57 def test_invalid(self):
58 with self.assertRaises(ValueError):
59 BlockStorageTypeFactory(None)
61 def test_register_invalid_name(self):
62 with self.assertRaises(ValueError):
63 BlockStorageTypeFactory.register_device(
64 's3', BlockStorageFile)
66 def test_register_invalid_type(self):
67 with self.assertRaises(TypeError):
68 BlockStorageTypeFactory.register_device(
71 class _TestBlockStorage(object):
77 def _read_storage(cls, storage):
78 with open(storage.storage_name, 'rb') as f:
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)
90 def _check_exists(cls, name):
91 return os.path.exists(name)
94 def _get_empty_existing(cls):
95 return os.path.join(thisdir,
100 def _get_dummy_noexist(cls):
101 fd, name = tempfile.mkstemp(dir=os.getcwd())
105 def _open_teststorage(self, **kwds):
106 kwds.update(self._type_kwds)
107 return self._type(self._testfname, **kwds)
109 def _reopen_storage(self, storage):
110 return self._type(storage.storage_name, **self._type_kwds)
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):
121 _remove_storage(cls._dummy_name) # pragma: no cover
124 cls._testfname = cls.__name__ + "_testfile.bin"
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,
135 for i in range(cls._block_count):
136 data = bytearray([i])*cls._block_size
137 cls._blocks.append(data)
140 def tearDownClass(cls):
141 cls._remove_storage(cls._testfname)
142 cls._remove_storage(cls._dummy_name)
144 def test_setup_fails(self):
145 self.assertEqual(self._check_exists(self._dummy_name), False)
146 with self.assertRaises(IOError):
148 self._get_empty_existing(),
152 self.assertEqual(self._check_exists(self._dummy_name), False)
153 with self.assertRaises(IOError):
155 self._get_empty_existing(),
158 ignore_existing=False,
160 self.assertEqual(self._check_exists(self._dummy_name), False)
161 with self.assertRaises(ValueError):
162 self._type.setup(self._dummy_name,
166 self.assertEqual(self._check_exists(self._dummy_name), False)
167 with self.assertRaises(ValueError):
168 self._type.setup(self._dummy_name,
172 self.assertEqual(self._check_exists(self._dummy_name), False)
173 with self.assertRaises(TypeError):
174 self._type.setup(self._dummy_name,
179 self.assertEqual(self._check_exists(self._dummy_name), False)
180 # TODO: The multiprocessing module is bad
181 # about handling exceptions raised on the
183 #with self.assertRaises(ValueError):
186 # self._type.setup(self._dummy_name,
191 #self.assertEqual(self._check_exists(self._dummy_name), False)
193 def test_setup(self):
194 fname = ".".join(self.id().split(".")[1:])
196 fname = os.path.join(thisdir, fname)
197 self._remove_storage(fname)
200 fsetup = self._type.setup(fname, bsize, bcount, **self._type_kwds)
202 flen = len(self._read_storage(fsetup))
205 self._type.compute_storage_size(bsize,
209 self._type.compute_storage_size(bsize,
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)
225 self.assertEqual(fsetup.storage_name, None)
226 self._remove_storage(fname)
228 def test_setup_withdata(self):
229 fname = ".".join(self.id().split(".")[1:])
231 fname = os.path.join(thisdir, fname)
232 self._remove_storage(fname)
235 header_data = bytes(bytearray([0,1,2]))
236 fsetup = self._type.setup(fname,
239 header_data=header_data,
243 flen = len(self._read_storage(fsetup))
246 self._type.compute_storage_size(bsize,
248 header_data=header_data))
249 self.assertTrue(len(header_data) > 0)
251 self._type.compute_storage_size(bsize,
253 self._type.compute_storage_size(bsize,
255 header_data=header_data),
259 self._type.compute_storage_size(bsize,
261 header_data=header_data,
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)
276 self.assertEqual(fsetup.storage_name, None)
278 self._remove_storage(fname)
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
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)
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)
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]))
316 self.assertEqual(f.bytes_sent, 0)
317 self.assertEqual(f.bytes_received,
318 self._block_count*self._block_size*4)
320 with self._open_teststorage() as f:
321 self.assertEqual(f.bytes_sent, 0)
322 self.assertEqual(f.bytes_received, 0)
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]))
329 self.assertEqual(f.bytes_sent, 0)
330 self.assertEqual(f.bytes_received,
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)
340 for i in xrange(self._block_count):
341 self.assertNotEqual(list(bytearray(f.read_block(i))),
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))),
348 for i, block in enumerate(self._blocks):
349 f.write_block(i, bytes(block))
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)
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)
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]))
379 self.assertEqual(f.bytes_sent, 0)
380 self.assertEqual(f.bytes_received,
381 (2*self._block_count+1)*self._block_size)
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)
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]))
406 self.assertEqual(f.bytes_sent, 0)
407 self.assertEqual(f.bytes_received,
408 (2*self._block_count+1)*self._block_size)
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)
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)),
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]))
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)
442 def test_update_header_data(self):
443 fname = ".".join(self.id().split(".")[1:])
445 fname = os.path.join(thisdir, fname)
446 self._remove_storage(fname)
449 header_data = bytes(bytearray([0,1,2]))
450 fsetup = self._type.setup(fname,
453 header_data=header_data,
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)
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:
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:
488 with self._open_teststorage(ignore_lock=True) as f1:
490 with self._open_teststorage(ignore_lock=True) as f:
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)
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]))
516 self.assertEqual(f.bytes_sent, 0)
517 self.assertEqual(f.bytes_received,
518 self._block_count*self._block_size*4)
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)
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]))
531 self.assertEqual(f.bytes_sent, 0)
532 self.assertEqual(f.bytes_received,
534 self.assertEqual(forig.bytes_sent, 0)
535 self.assertEqual(forig.bytes_received, 0)
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)
549 for i in xrange(self._block_count):
550 self.assertNotEqual(list(bytearray(f.read_block(i))),
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))),
557 for i, block in enumerate(self._blocks):
558 f.write_block(i, bytes(block))
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)
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)
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]))
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)
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)
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]))
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)
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)
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)),
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]))
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)
674 class TestBlockStorageFile(_TestBlockStorage,
676 _type = BlockStorageFile
679 class TestBlockStorageFileNoThreadPool(_TestBlockStorage,
681 _type = BlockStorageFile
682 _type_kwds = {'threadpool_size': 0}
684 class TestBlockStorageFileThreadPool(_TestBlockStorage,
686 _type = BlockStorageFile
687 _type_kwds = {'threadpool_size': 1}
689 class TestBlockStorageMMap(_TestBlockStorage,
691 _type = BlockStorageMMap
694 class _TestBlockStorageRAM(_TestBlockStorage):
697 def _read_storage(cls, storage):
701 def _remove_storage(cls, name):
705 def _check_exists(cls, name):
708 def _open_teststorage(self, **kwds):
709 kwds.update(self._type_kwds)
710 return self._type(self._original_f.data, **kwds)
712 def _reopen_storage(self, storage):
713 return self._type(storage.data, **self._type_kwds)
716 # Override some of the test methods
719 def test_setup_fails(self):
720 with self.assertRaises(ValueError):
721 self._type.setup(self._dummy_name,
725 with self.assertRaises(ValueError):
726 self._type.setup(self._dummy_name,
730 with self.assertRaises(TypeError):
731 self._type.setup(self._dummy_name,
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
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)
759 def test_tofile_fromfile_fileobj(self):
761 self._original_f.tofile(out1)
763 self.assertEqual(len(self._original_f.data) > 0, True)
764 self.assertEqual(self._original_f.data, out1.read())
766 in1 = self._type.fromfile(out1)
767 self.assertNotEqual(self._original_f.data, in1.data)
770 self.assertNotEqual(self._original_f.data, in1.data)
772 self.assertEqual(self._original_f.data, in1.data)
774 with self.assertRaises(IOError):
775 with self._type.fromfile(out2) as in2:
776 pass # pragma: no cover
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)
784 def test_tofile_fromfile_filename(self):
787 fd, out = tempfile.mkstemp()
791 with open(name, 'rb') as f:
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)
802 self.assertNotEqual(self._original_f.data, in1.data)
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)
814 class TestBlockStorageRAM(_TestBlockStorageRAM,
816 _type = BlockStorageRAM
819 class _dummy_sftp_file(object):
820 def __init__(self, *args, **kwds):
821 self._f = open(*args, **kwds)
824 def __exit__(self, *args):
826 def readv(self, chunks):
828 for offset, size in chunks:
830 data.append(self._f.read(size))
832 def __getattr__(self, key):
833 return getattr(self._f, key)
834 def set_pipelined(self):
837 class dummy_sftp(object):
841 def open(*args, **kwds):
842 return _dummy_sftp_file(*args, **kwds)
847 class dummy_sshclient(object):
852 class TestBlockStorageSFTP(_TestBlockStorage,
854 _type = BlockStorageSFTP
855 _type_kwds = {'sshclient': dummy_sshclient}
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,
866 self.assertEqual(self._check_exists(self._dummy_name), False)
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
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)
887 class _TestBlockStorageS3Mock(_TestBlockStorage):
888 _type = BlockStorageS3
892 def _read_storage(cls, storage):
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())
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)
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,
933 self.assertEqual(self._check_exists(self._dummy_name), False)
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,
940 **self._type_kwds) as f:
942 self.assertEqual(self._check_exists(self._dummy_name), True)
943 with self.assertRaises(IOError):
944 with self._type.setup(self._dummy_name,
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,
953 ignore_existing=True,
954 **self._type_kwds) as f:
956 self.assertEqual(self._check_exists(self._dummy_name), True)
957 self._remove_storage(self._dummy_name)
959 class TestBlockStorageS3Mock(_TestBlockStorageS3Mock,
961 _type_kwds = {'s3_wrapper': MockBoto3S3Wrapper,
964 class TestBlockStorageS3MockNoThreadPool(_TestBlockStorageS3Mock,
966 _type_kwds = {'s3_wrapper': MockBoto3S3Wrapper,
968 'threadpool_size': 0}
970 class TestBlockStorageS3MockThreadPool(_TestBlockStorageS3Mock,
972 _type_kwds = {'s3_wrapper': MockBoto3S3Wrapper,
974 'threadpool_size': 4}
976 @unittest2.skipIf((os.environ.get('PYORAM_AWS_TEST_BUCKET') is None) or \
978 "No PYORAM_AWS_TEST_BUCKET defined in environment or "
979 "boto3 is not available")
980 class TestBlockStorageS3(_TestBlockStorage,
982 _type = BlockStorageS3
983 _type_kwds = {'bucket_name': os.environ.get('PYORAM_AWS_TEST_BUCKET')}
986 def _read_storage(cls, storage):
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)))
999 def _remove_storage(cls, name):
1000 Boto3S3Wrapper(cls._type_kwds['bucket_name']).clear(name)
1003 def _check_exists(cls, name):
1004 return Boto3S3Wrapper(cls._type_kwds['bucket_name']).exists(name)
1007 def _get_empty_existing(cls):
1008 return "exists.empty"
1011 def _get_dummy_noexist(cls):
1012 s3 = Boto3S3Wrapper(cls._type_kwds['bucket_name'])
1013 fd, name = tempfile.mkstemp(dir=os.getcwd())
1016 while s3.exists(name):
1017 fd, name = tempfile.mkstemp(dir=os.getcwd())
1022 if __name__ == "__main__":
1023 unittest2.main() # pragma: no cover