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_remove(self):
"""Test the __contains__ method"""
data = "abcdef"
to_remove = "bcd"
urn = Urn(data, replace=False)
# Draw one sample, then verify that element is not still in the urn
for element in to_remove:
assert element in urn
urn.remove(element)
assert element not in urn