def test_get_bloom_filter(self):
s = server.StorageServer(test_dir)
b = BloomFilter(20, 0.01, self.bloom_path) # create bloom filter
b.update(l2)
b64 = s.get_bloom_filter()
b = BloomFilter.from_base64(f"{test_dir}bloom-test.bloom", b64)
for e in l2:
self.assertIn(e, b)
for e in l3:
self.assertNotIn(e, b)
def test_store_record(self):
s = server.StorageServer(test_dir)
BloomFilter(20, 0.1, self.bloom_path) # create bloom filter
self.assertNotIn('a', s.bloom)
with patch("lib.storage_server_backend.db") as db:
s.store_record('a', 'record', 'owner')
# Check bloom filter
self.assertIn('a', s.bloom)
# check db
db.session.add.assert_called_once()
def test_batch_store_records_bloom(self):
s = server.StorageServer(test_dir)
BloomFilter(20, 0.1, self.bloom_path) # create bloom filter
b = s.bloom
for e in l2:
self.assertNotIn(e, b)
s.batch_store_records_bloom(l1)
# Check bloom filter
for e in l2:
self.assertIn(e, b)
def test_bloom(self):
s = server.StorageServer(test_dir)
self.assertFalse(os.path.exists(self.bloom_path))
with patch.object(s, "_initialize_bloom_filter") as m:
b = s.bloom
m.assert_called_once()
b = BloomFilter(2, .1, self.bloom_path)
b.add(5)
c = s.bloom
m.assert_called_once() # No second call
self.assertIn(5, c)
def test_batch_get_record(self):
server.db.init_app(mock_app)
with mock_app.test_request_context(), \
patch.object(server.StorageServer, "bloom", new_callable=Mock()):
server.db.create_all()
s = server.StorageServer(test_dir)
s.batch_store_records_db(l1)
s.batch_store_records_bloom(l1)
res = s.batch_get_records(l2[::2], "client")
self.assertEqual(
[(h, r) for (h, r, o) in l1[::2]],
res
)
def test_get_all_record_psi_hashes(self):
records = []
correct = []
for i in range(10):
r = Record([1, 2, 3, 4, 5])
r.set_hash_key(b'fake_key')
m = Mock()
m.hash = helpers.to_base64(r.get_long_hash())
records.append(m)
correct.append(r.get_psi_index())
with patch("lib.storage_server_backend.StoredRecord") as c:
c.query.all.return_value = records
s = server.StorageServer()
res = s.get_all_record_psi_hashes()
self.assertEqual(correct, res)
def test_initialize_bloom_filter(self, m):
m.query.all.return_value = []
s = server.StorageServer(test_dir)
# Update with contents from DB
mock_records = []
for l in l2:
mo = Mock()
mo.hash = l
mock_records.append(mo)
m.query.all.return_value = mock_records
s._initialize_bloom_filter()
b = s.bloom
for e in l2:
self.assertIn(e, b)
for e in l3:
self.assertNotIn(e, b)
def test_offer_psi(self):
port = 5555
s = server.StorageServer()
set = list(range(20))
m = Mock()
with mock.patch("lib.storage_server_backend.PyPSISender",
return_value=m):
with mock.patch.object(server.StorageServer,
"get_all_record_psi_hashes",
return_value=set):
s.offer_psi(22, port)
self.assertEqual(config.PSI_SCHEME, m.execute.call_args[0][0])
self.assertEqual(set, m.execute.call_args[0][1])
set = list(range(100))
mock.patch.object(s, "get_all_record_psi_hashes", return_value=set)
with self.assertRaises(RuntimeError):
with mock.patch.object(server.StorageServer,
"get_all_record_psi_hashes",
return_value=set):
s.offer_psi(22, port=port)
def test_init(self):
# Is Directory created?
os.rmdir(test_dir)
self.assertFalse(os.path.exists(test_dir))
server.StorageServer(test_dir)
self.assertTrue(os.path.exists(test_dir))