diff --git a/GNNlib/test/Project.toml b/GNNlib/test/Project.toml
new file mode 100644
index 000000000..d7161d336
--- /dev/null
+++ b/GNNlib/test/Project.toml
@@ -0,0 +1,12 @@
+[deps]
+GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
+GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
diff --git a/GNNlib/test/msgpass_tests.jl b/GNNlib/test/msgpass.jl
similarity index 98%
rename from GNNlib/test/msgpass_tests.jl
rename to GNNlib/test/msgpass.jl
index 60d13fcb0..186b4dc78 100644
--- a/GNNlib/test/msgpass_tests.jl
+++ b/GNNlib/test/msgpass.jl
@@ -1,4 +1,5 @@
-@testitem "msgpass" setup=[SharedTestSetup] begin
+@testitem "msgpass" setup=[TestModuleNNlib] begin
+    using .TestModuleNNlib
     #TODO test all graph types
     GRAPH_T = :coo
     in_channel = 10
@@ -137,4 +138,4 @@
         end
     end
 
-end
\ No newline at end of file
+end
diff --git a/GNNlib/test/runtests.jl b/GNNlib/test/runtests.jl
index e4c4512b4..d420f5a6c 100644
--- a/GNNlib/test/runtests.jl
+++ b/GNNlib/test/runtests.jl
@@ -1,6 +1,28 @@
-using GNNlib
-using Test
-using ReTestItems
-using Random, Statistics
+using TestItemRunner
 
-runtests(GNNlib)
+## See https://www.julia-vscode.org/docs/stable/userguide/testitems/
+## for how to run the tests within VS Code.
+## See test_module.jl for the test infrastructure.
+
+## Uncomment below and in test_module.jl to change the default test settings
+# ENV["GNN_TEST_CPU"] = "false"
+# ENV["GNN_TEST_CUDA"] = "true"
+# ENV["GNN_TEST_AMDGPU"] = "true"
+# ENV["GNN_TEST_Metal"] = "true"
+
+# The only available tag at the moment is :gpu
+# Tests not tagged with :gpu are considered to be CPU tests
+# Tests tagged with :gpu should run on all GPU backends
+
+if get(ENV, "GNN_TEST_CPU", "true") == "true"
+    @run_package_tests filter = ti -> :gpu ∉ ti.tags
+end
+if get(ENV, "GNN_TEST_CUDA", "false") == "true"
+    @run_package_tests filter = ti -> :gpu ∈ ti.tags
+end
+if get(ENV, "GNN_TEST_AMDGPU", "false") == "true"
+    @run_package_tests filter = ti -> :gpu ∈ ti.tags
+end
+if get(ENV, "GNN_TEST_Metal", "false") == "true"
+    @run_package_tests filter = ti -> :gpu ∈ ti.tags
+end
diff --git a/GNNlib/test/shared_testsetup.jl b/GNNlib/test/test_module.jl
similarity index 80%
rename from GNNlib/test/shared_testsetup.jl
rename to GNNlib/test/test_module.jl
index 106db5159..a2229ef59 100644
--- a/GNNlib/test/shared_testsetup.jl
+++ b/GNNlib/test/test_module.jl
@@ -1,4 +1,4 @@
-@testsetup module SharedTestSetup
+@testmodule TestModuleNNlib begin
 
 import Reexport: @reexport
 
@@ -9,4 +9,4 @@ import Reexport: @reexport
 @reexport using SparseArrays
 @reexport using Test, Random, Statistics
 
-end
\ No newline at end of file
+end # module
\ No newline at end of file
diff --git a/GNNlib/test/utils_tests.jl b/GNNlib/test/utils.jl
similarity index 100%
rename from GNNlib/test/utils_tests.jl
rename to GNNlib/test/utils.jl