function test_PartitionTS()
    d = 2
    lims = repeat([-1.0 1.0], d, 1)
    traj = Trajectory(d)
    retrain = [10 * 2^i for i in 1:5]
    pts = PartitionTS(1, retrain, traj, lims, 3, 100, p -> 4 * p, 2, false)
    f(x) = 5 * x[1]^3* x[2]^4 - 5 * x[2]^2 + 7 * x[1] * x[2]
    n = retrain[end]
    X, r = AdaptivePartition.gaussian_data(f, d, n, 1.0, lims)
    ContextualBandits.add_memory!(traj, n)
    for i in 1:n
        ContextualBandits.extend!(traj, X[i, :], 1, r[i], 0.0)
        update!(pts, X[i, :], 1, r[i])
    end
    m = partition_search(X, r, lims, 3, 100, p -> 4 * p, 0, false)
    @test m.scale ≈ pts.arms[1].scale
    @test choose([0.0, 0.0], pts) == 1
end

test_PartitionTS()