From 664f91bfab921a07d692b35e9d605510d0d8b891 Mon Sep 17 00:00:00 2001 From: Jack Robison Date: Fri, 15 Mar 2019 12:44:41 -0400 Subject: [PATCH] add lbrynet.dht.protocol.distance unit tests --- lbrynet/dht/protocol/distance.py | 4 +++- tests/unit/dht/protocol/test_distance.py | 13 +++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) create mode 100644 tests/unit/dht/protocol/test_distance.py diff --git a/lbrynet/dht/protocol/distance.py b/lbrynet/dht/protocol/distance.py index 516dfabc4..2b2577fa3 100644 --- a/lbrynet/dht/protocol/distance.py +++ b/lbrynet/dht/protocol/distance.py @@ -10,11 +10,13 @@ class Distance: def __init__(self, key: bytes): if len(key) != constants.hash_length: - raise ValueError("invalid key length: %i" % len(key)) + raise ValueError(f"invalid key length: {len(key)}") self.key = key self.val_key_one = int.from_bytes(key, 'big') def __call__(self, key_two: bytes) -> int: + if len(key_two) != constants.hash_length: + raise ValueError(f"invalid length of key to compare: {len(key_two)}") val_key_two = int.from_bytes(key_two, 'big') return self.val_key_one ^ val_key_two diff --git a/tests/unit/dht/protocol/test_distance.py b/tests/unit/dht/protocol/test_distance.py new file mode 100644 index 000000000..829bcf3f9 --- /dev/null +++ b/tests/unit/dht/protocol/test_distance.py @@ -0,0 +1,13 @@ +import unittest +from lbrynet.dht.protocol.distance import Distance + + +class DistanceTests(unittest.TestCase): + def test_invalid_key_length(self): + self.assertRaises(ValueError, Distance, b'1' * 47) + self.assertRaises(ValueError, Distance, b'1' * 49) + self.assertRaises(ValueError, Distance, b'') + + self.assertRaises(ValueError, Distance(b'0' * 48), b'1' * 47) + self.assertRaises(ValueError, Distance(b'0' * 48), b'1' * 49) + self.assertRaises(ValueError, Distance(b'0' * 48), b'')