diff --git a/src/lib/GraphsExtensions/src/neighbors.jl b/src/lib/GraphsExtensions/src/neighbors.jl
index be6fdda..cb62a29 100644
--- a/src/lib/GraphsExtensions/src/neighbors.jl
+++ b/src/lib/GraphsExtensions/src/neighbors.jl
@@ -1,10 +1,8 @@
 using Graphs: AbstractGraph, neighborhood_dists
 
-function vertices_at_distance(g::AbstractGraph, vertex, n::Int)
-  neighborhood = [first(v) for v in neighborhood_dists(g, vertex, n)]
-  closer_neighborhood = [first(v) for v in neighborhood_dists(g, vertex, n - 1)]
-  iszero(n) && return neighborhood
-  return setdiff(neighborhood, closer_neighborhood)
+function vertices_at_distance(g::AbstractGraph, vertex, distance::Int)
+  n_dists = neighborhood_dists(g, vertex, distance)
+  return map(first, filter(==(distance) ∘ last, n_dists))
 end
 
 next_nearest_neighbors(g::AbstractGraph, v) = vertices_at_distance(g, v, 2)
diff --git a/src/lib/GraphsExtensions/test/runtests.jl b/src/lib/GraphsExtensions/test/runtests.jl
index b645b64..f6533a4 100644
--- a/src/lib/GraphsExtensions/test/runtests.jl
+++ b/src/lib/GraphsExtensions/test/runtests.jl
@@ -69,6 +69,7 @@ using NamedGraphs.GraphsExtensions:
   is_self_loop,
   leaf_vertices,
   minimum_distance_to_leaves,
+  next_nearest_neighbors,
   non_leaf_edges,
   outdegrees,
   permute_vertices,
@@ -81,7 +82,8 @@ using NamedGraphs.GraphsExtensions:
   tree_graph_node,
   undirected_graph,
   undirected_graph_type,
-  vertextype
+  vertextype,
+  vertices_at_distance
 using Test: @test, @test_broken, @test_throws, @testset
 
 # TODO: Still need to test:
@@ -579,5 +581,12 @@ using Test: @test, @test_broken, @test_throws, @testset
   g′ = path_graph(4)
   rem_edges!(g′, [2 => 3, 3 => 4])
   @test g′ == g
+
+  #vertices at distance
+  L = 10
+  g = path_graph(L)
+  @test only(vertices_at_distance(g, 1, L - 1)) == L
+  @test only(next_nearest_neighbors(g, 1)) == 3
+  @test length(vertices_at_distance(g, 5, 3)) == 2
 end
 end
diff --git a/test/test_graphsextensions.jl b/test/test_graphsextensions.jl
deleted file mode 100644
index 6a3607d..0000000
--- a/test/test_graphsextensions.jl
+++ /dev/null
@@ -1,24 +0,0 @@
-@eval module $(gensym())
-using Graphs: AbstractGraph, neighborhood_dists
-using NamedGraphs.NamedGraphGenerators: named_grid
-using NamedGraphs.GraphsExtensions: next_nearest_neighbors, vertices_at_distance
-using Test: @test, @testset
-
-#TODO: Add tests for other graphs extensions
-@testset "GraphsExtensions" begin
-  @testset "Test vertices at distance" begin
-    L = 10
-    g = named_grid((L, 1))
-    vstart = (1, 1)
-    @test only(vertices_at_distance(g, vstart, L - 1)) == (L, 1)
-    @test only(next_nearest_neighbors(g, vstart)) == (3, 1)
-
-    L = 9
-    g = named_grid((L, L))
-    v_middle = (ceil(Int64, L / 2), ceil(Int64, L / 2))
-    corners = [(L, 1), (1, L), (L, L), (1, 1)]
-    @test length(next_nearest_neighbors(g, v_middle)) == 8
-    @test issetequal(vertices_at_distance(g, v_middle, L - 1), corners)
-  end
-end
-end