def test_list_extension(self):
"""Test the extend method"""
data = "abcdef"
urn = Urn(data, replace=False)
assert urn.size() == 6
urn.extend("xy")
assert urn.size() == 8
def test_urn_length_without_replacement(self, weights):
"""Test that urn.size() works as expected.
"""
data = [1, 2, 3]
urn = Urn(data, replace=False, weights=weights)
assert urn.size() == 3
next(urn)
assert urn.size() == 2
list(urn) # Exhaust the iterator
assert urn.size() == 0
def test_remove_weighted(self, replace):
data = "abcdef"
weights = [1, 2, 3, 4, 5, 6]
to_remove = "bde"
urn = Urn(data, replace, weights)
for element in to_remove:
assert element in urn
urn.remove(element)
assert element not in urn
assert urn.size() == 3 or urn.size() == float("inf")
def test_add_weighted(self, replace):
data = "acf"
weights = [1, 2, 3]
more_data = "bde"
more_weights = [4, 5, 6]
urn = Urn(data, replace, weights)
for element, weight in zip(more_data, more_weights):
assert element not in urn
urn.add(element, weight)
assert element in urn
assert urn.size() == 6 or urn.size() == float("inf")
def test_urn_length_with_replacement(self, weights):
"""Test that urn.size() works as expected.
"""
data = [1, 2, 3]
urn = Urn(data, replace=True, weights=weights)
assert urn.size() == float("inf")
# Drawing a sample does not change the size
next(urn)
assert urn.size() == float("inf")
# Drawing 50 samples does not change the size
list(itertools.islice(urn, 50))
assert urn.size() == float("inf")
def test_urn_bool(self, weights):
"""Test that bool(urn) works as expected.
"""
data = [1, 2, 3]
urn = Urn(data, replace=False, weights=weights)
assert urn
assert bool(urn)
while urn:
next(urn)
assert urn.size() == 0
assert not urn
assert not bool(urn)