function test_KNNUCB()
    d = 2
    A = 2
    traj = Trajectory(d)
    pol = KNNUCB(A, traj, t-> 1.0, 1.0)
    ContextualBandits.add_memory!(traj, 3)

    X = [
            0.0 0.1; # a = 1
            0.0 0.1; # a = 2
            0.0 0.2 # a = 1
        ]
    for i in 1:A
        a = choose(X[i, :], pol)
        @test a == i
        ContextualBandits.extend!(traj, X[i, :], a, Float64(i), 0.0)
        ContextualBandits.update!(pol, X[i, :], a, Float64(i))
    end


    ContextualBandits.extend!(traj, X[3, :], 1, 2.0, 0.0)
    ContextualBandits.update!(pol, X[3, :], 1, 2.0)
    x0 = [0.0, 0.0]
    dst, r = ContextualBandits.sort_nearest(pol, x0, 1)
    @test dst == [0.1, 0.2]
    @test r == [1.0, 2.0]
end

test_KNNUCB()

