use the new helpers for realistic files

Signed-off-by: Alfredo Deza <alfredo@deza.pe>
This commit is contained in:
Alfredo Deza 2014-12-12 13:53:19 -05:00
parent b4794edbb3
commit cad6547eb0

View file

@ -21,6 +21,7 @@ def generate_file_contents(size):
should remove the last 40 chars from the blob to retrieve the original hash should remove the last 40 chars from the blob to retrieve the original hash
and binary so that validity can be proved. and binary so that validity can be proved.
""" """
size = int(size)
contents = os.urandom(size) contents = os.urandom(size)
content_hash = hashlib.sha1(contents).hexdigest() content_hash = hashlib.sha1(contents).hexdigest()
return contents + content_hash return contents + content_hash
@ -28,22 +29,29 @@ def generate_file_contents(size):
class FileValidator(object): class FileValidator(object):
def __init__(self, f): def __init__(self, f=None):
self._file = tempfile.SpooledTemporaryFile() self._file = tempfile.SpooledTemporaryFile()
f.seek(0) self.original_hash = None
shutil.copyfileobj(f, self._file) self.new_hash = None
self.seek(0) if f:
f.seek(0)
shutil.copyfileobj(f, self._file)
def valid(self): def valid(self):
""" """
Returns True if this file looks valid. The file is valid if the end Returns True if this file looks valid. The file is valid if the end
of the file has the md5 digest for the first part of the file. of the file has the md5 digest for the first part of the file.
""" """
contents = self._file.read()
self._file.seek(0) self._file.seek(0)
original_hash, binary = contents[-40:], contents[:-40] contents = self._file.read()
new_hash = hashlib.sha1(binary).hexdigest() self.original_hash, binary = contents[-40:], contents[:-40]
return new_hash == original_hash self.new_hash = hashlib.sha1(binary).hexdigest()
if not self.new_hash == self.original_hash:
print 'original hash: ', self.original_hash
print 'new hash: ', self.new_hash
print 'size: ', self._file.tell()
return False
return True
# XXX not sure if we need all of these # XXX not sure if we need all of these
def seek(self, offset, whence=os.SEEK_SET): def seek(self, offset, whence=os.SEEK_SET):
@ -55,6 +63,10 @@ class FileValidator(object):
def read(self, size=-1): def read(self, size=-1):
return self._file.read(size) return self._file.read(size)
def write(self, data):
self._file.write(data)
self._file.seek(0)
class RandomContentFile(object): class RandomContentFile(object):
def __init__(self, size, seed): def __init__(self, size, seed):
@ -228,6 +240,7 @@ def files(mean, stddev, seed=None):
break break
yield RandomContentFile(size=size, seed=rand.getrandbits(32)) yield RandomContentFile(size=size, seed=rand.getrandbits(32))
def files2(mean, stddev, seed=None, numfiles=10): def files2(mean, stddev, seed=None, numfiles=10):
""" """
Yields file objects with effectively random contents, where the Yields file objects with effectively random contents, where the
@ -238,17 +251,17 @@ def files2(mean, stddev, seed=None, numfiles=10):
stores `numfiles` files and yields them in a loop. stores `numfiles` files and yields them in a loop.
""" """
# pre-compute all the files (and save with TemporaryFiles) # pre-compute all the files (and save with TemporaryFiles)
rand_files = files(mean, stddev, seed)
fs = [] fs = []
for _ in xrange(numfiles): for _ in xrange(numfiles):
f = next(rand_files)
t = tempfile.SpooledTemporaryFile() t = tempfile.SpooledTemporaryFile()
shutil.copyfileobj(f, t) t.write(generate_file_contents(random.normalvariate(mean, stddev)))
t.seek(0)
fs.append(t) fs.append(t)
while True: while True:
for f in fs: for f in fs:
yield PrecomputedContentFile(f) yield f
def names(mean, stddev, charset=None, seed=None): def names(mean, stddev, charset=None, seed=None):
""" """