fix CompressionSpec validator and tests

(cherry picked from commit 0b40e038ad)
This commit is contained in:
Thomas Waldmann 2023-04-12 01:45:46 +02:00 committed by Sasha Boginsky
parent de199c23ab
commit 3e9668ecbc
2 changed files with 15 additions and 13 deletions

View file

@ -15,6 +15,7 @@ which compressor has been used to compress the data and dispatch to the correct
decompressor.
"""
from argparse import ArgumentTypeError
import random
from struct import Struct
import zlib
@ -551,7 +552,7 @@ class CompressionSpec:
values = s.split(',')
count = len(values)
if count < 1:
raise ValueError
raise ArgumentTypeError("not enough arguments")
# --compression algo[,level]
self.name = values[0]
if self.name in ('none', 'lz4', ):
@ -562,9 +563,9 @@ class CompressionSpec:
elif count == 2:
level = int(values[1])
if not 0 <= level <= 9:
raise ValueError
raise ArgumentTypeError("level must be >= 0 and <= 9")
else:
raise ValueError
raise ArgumentTypeError("too many arguments")
self.level = level
elif self.name in ('zstd', ):
if count < 2:
@ -572,28 +573,28 @@ class CompressionSpec:
elif count == 2:
level = int(values[1])
if not 1 <= level <= 22:
raise ValueError
raise ArgumentTypeError("level must be >= 1 and <= 22")
else:
raise ValueError
raise ArgumentTypeError("too many arguments")
self.level = level
elif self.name == 'auto':
if 2 <= count <= 3:
compression = ','.join(values[1:])
else:
raise ValueError
raise ArgumentTypeError("bad arguments")
self.inner = CompressionSpec(compression)
elif self.name == 'obfuscate':
if 3 <= count <= 5:
level = int(values[1])
if not ((1 <= level <= 6) or (110 <= level <= 123)):
raise ValueError
raise ArgumentTypeError("level must be >= 1 and <= 6 or >= 110 and <= 123")
self.level = level
compression = ','.join(values[2:])
else:
raise ValueError
raise ArgumentTypeError("bad arguments")
self.inner = CompressionSpec(compression)
else:
raise ValueError
raise ArgumentTypeError("unsupported compression type")
@property
def compressor(self):

View file

@ -1,3 +1,4 @@
import argparse
import os
import zlib
try:
@ -198,7 +199,7 @@ def test_obfuscate():
def test_compression_specs():
with pytest.raises(ValueError):
with pytest.raises(argparse.ArgumentTypeError):
CompressionSpec('')
assert isinstance(CompressionSpec('none').compressor, CNONE)
@ -213,7 +214,7 @@ def test_compression_specs():
zlib = CompressionSpec('zlib,9').compressor
assert isinstance(zlib, ZLIB)
assert zlib.level == 9
with pytest.raises(ValueError):
with pytest.raises(argparse.ArgumentTypeError):
CompressionSpec('zlib,9,invalid')
lzma = CompressionSpec('lzma').compressor
@ -236,7 +237,7 @@ def test_compression_specs():
assert isinstance(zstd, ZSTD)
assert zstd.level == 22
with pytest.raises(ValueError):
with pytest.raises(argparse.ArgumentTypeError):
CompressionSpec('lzma,9,invalid')
with pytest.raises(ValueError):
with pytest.raises(argparse.ArgumentTypeError):
CompressionSpec('invalid')