test_wheelfile.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. import wheel.install
  2. import hashlib
  3. try:
  4. from StringIO import StringIO
  5. except ImportError:
  6. from io import BytesIO as StringIO
  7. import zipfile
  8. import pytest
  9. def test_verifying_zipfile():
  10. if not hasattr(zipfile.ZipExtFile, '_update_crc'):
  11. pytest.skip('No ZIP verification. Missing ZipExtFile._update_crc.')
  12. sio = StringIO()
  13. zf = zipfile.ZipFile(sio, 'w')
  14. zf.writestr("one", b"first file")
  15. zf.writestr("two", b"second file")
  16. zf.writestr("three", b"third file")
  17. zf.close()
  18. # In default mode, VerifyingZipFile checks the hash of any read file
  19. # mentioned with set_expected_hash(). Files not mentioned with
  20. # set_expected_hash() are not checked.
  21. vzf = wheel.install.VerifyingZipFile(sio, 'r')
  22. vzf.set_expected_hash("one", hashlib.sha256(b"first file").digest())
  23. vzf.set_expected_hash("three", "blurble")
  24. vzf.open("one").read()
  25. vzf.open("two").read()
  26. try:
  27. vzf.open("three").read()
  28. except wheel.install.BadWheelFile:
  29. pass
  30. else:
  31. raise Exception("expected exception 'BadWheelFile()'")
  32. # In strict mode, VerifyingZipFile requires every read file to be
  33. # mentioned with set_expected_hash().
  34. vzf.strict = True
  35. try:
  36. vzf.open("two").read()
  37. except wheel.install.BadWheelFile:
  38. pass
  39. else:
  40. raise Exception("expected exception 'BadWheelFile()'")
  41. vzf.set_expected_hash("two", None)
  42. vzf.open("two").read()
  43. def test_pop_zipfile():
  44. sio = StringIO()
  45. zf = wheel.install.VerifyingZipFile(sio, 'w')
  46. zf.writestr("one", b"first file")
  47. zf.writestr("two", b"second file")
  48. zf.close()
  49. try:
  50. zf.pop()
  51. except RuntimeError:
  52. pass # already closed
  53. else:
  54. raise Exception("expected RuntimeError")
  55. zf = wheel.install.VerifyingZipFile(sio, 'a')
  56. zf.pop()
  57. zf.close()
  58. zf = wheel.install.VerifyingZipFile(sio, 'r')
  59. assert len(zf.infolist()) == 1