lbry-sdk/lbrynet/tests/unit/lbryfilemanager/test_EncryptedFileCreator.py

69 lines
2.6 KiB
Python

# -*- coding: utf-8 -*-
from Crypto.Cipher import AES
import mock
from twisted.trial import unittest
from twisted.internet import defer
from lbrynet.core import BlobManager
from lbrynet.core import Session
from lbrynet.core.server import DHTHashAnnouncer
from lbrynet.file_manager import EncryptedFileCreator
from lbrynet.file_manager import EncryptedFileManager
from lbrynet.tests import mocks
from lbrynet.tests.util import mk_db_and_blob_dir, rm_db_and_blob_dir
MB = 2**20
def iv_generator():
while True:
yield '3' * AES.block_size
class CreateEncryptedFileTest(unittest.TestCase):
timeout = 5
def setUp(self):
mocks.mock_conf_settings(self)
self.tmp_db_dir, self.tmp_blob_dir = mk_db_and_blob_dir()
@defer.inlineCallbacks
def tearDown(self):
yield self.blob_manager.stop()
rm_db_and_blob_dir(self.tmp_db_dir, self.tmp_blob_dir)
@defer.inlineCallbacks
def create_file(self, filename):
session = mock.Mock(spec=Session.Session)(None, None)
hash_announcer = DHTHashAnnouncer.DHTHashAnnouncer(None, None)
self.blob_manager = BlobManager.DiskBlobManager(
hash_announcer, self.tmp_blob_dir, self.tmp_db_dir)
session.blob_manager = self.blob_manager
yield session.blob_manager.setup()
session.db_dir = self.tmp_db_dir
manager = mock.Mock(spec=EncryptedFileManager.EncryptedFileManager)()
handle = mocks.GenFile(3*MB, '1')
key = '2'*AES.block_size
out = yield EncryptedFileCreator.create_lbry_file(
session, manager, filename, handle, key, iv_generator())
defer.returnValue(out)
@defer.inlineCallbacks
def test_can_create_file(self):
expected_stream_hash = ('41e6b247d923d191b154fb6f1b8529d6ddd6a73d65c357b1acb7'
'42dd83151fb66393a7709e9f346260a4f4db6de10c25')
filename = 'test.file'
stream_hash = yield self.create_file(filename)
self.assertEqual(expected_stream_hash, stream_hash)
blobs = yield self.blob_manager.get_all_verified_blobs()
self.assertEqual(2, len(blobs))
num_should_announce_blobs = yield self.blob_manager.count_should_announce_blobs()
self.assertEqual(1, num_should_announce_blobs)
@defer.inlineCallbacks
def test_can_create_file_with_unicode_filename(self):
expected_stream_hash = ('d1da4258f3ce12edb91d7e8e160d091d3ab1432c2e55a6352dce0'
'2fd5adb86fe144e93e110075b5865fff8617776c6c0')
filename = u'☃.file'
stream_hash = yield self.create_file(filename)
self.assertEqual(expected_stream_hash, stream_hash)