def test_get_all(self): nn = KDTree_3Float() o1 = object() nn.add(((1,1,0), id(o1))) o2 = object() nn.add(((10,10,0), id(o2))) o3 = object() nn.add(((11,11,0), id(o3))) self.assertEqual([((1,1,0), id(o1)), ((10,10,0), id(o2)), ((11,11,0), id(o3))], nn.get_all()) self.assertEqual(3, len(nn)) nn.remove(((10,10,0), id(o2))) self.assertEqual(2, len(nn)) self.assertEqual([((1,1,0), id(o1)), ((11,11,0), id(o3))], nn.get_all())
def test_nearest(self): nn = KDTree_3Float() nn_id = {} o1 = object() nn.add(((1,1,0), id(o1))) nn_id[id(o1)] = o1 o2 = object() nn.add(((10,10,0), id(o2))) nn_id[id(o2)] = o2 o3 = object() nn.add(((4.1, 4.1,0), id(o3))) nn_id[id(o3)] = o3 expected = o3 actual = nn.find_nearest((2.9,2.9,0))[1] self.assertTrue(expected==nn_id[actual], "%s != %s"%(str(expected), str(nn_id[actual]))) expected = o3 actual = nn.find_nearest((6, 6,0))[1] self.assertTrue(expected==nn_id[actual], "%s != %s"%(str(expected), str(nn_id[actual])))
def test_remove(self): class C: def __init__(self, i): self.i = i self.next = None nn = KDTree_3Float() k1, o1 = (1.1, 1.1, 0), C(7) self.assertFalse(nn.remove((k1, id(o1))), "This cannot be removed!") nn.add((k1, id(o1))) k2, o2 = (1.1, 1.1, 0), C(7) nn.add((k2, id(o2))) self.assertEqual(2, nn.size()) self.assertTrue(nn.remove((k2, id(o2)))) self.assertEqual(1, nn.size()) self.assertFalse(nn.remove((k2, id(o2)))) self.assertEqual(1, nn.size()) nearest = nn.find_nearest(k1) self.assertTrue(nearest[1] == id(o1), "%s != %s" % (nearest[1], o1))
def test_empty(self): nn = KDTree_3Float() self.assertEqual(0, nn.size()) actual = nn.find_nearest((2, 3, 0)) self.assertTrue(None == actual, "%s != %s" % (str(None), str(actual)))