diff --git a/.ci/scripts/unittest-buck2.sh b/.ci/scripts/unittest-buck2.sh
index 09275f29ae9..4b614e32320 100755
--- a/.ci/scripts/unittest-buck2.sh
+++ b/.ci/scripts/unittest-buck2.sh
@@ -8,11 +8,17 @@ set -eux
 
 # TODO: expand this to //...
 # TODO: can't query cadence & vulkan backends
+# TODO: can't query //kernels/prim_ops because of a cpp_unittest and
+# broken code in shim to read oss.folly_cxx_tests. Sending fix but it
+# needs to propagate and we need a submodule update.
 buck2 query "//backends/apple/... + //backends/example/... + \
 //backends/mediatek/... + //backends/test/... + //backends/transforms/... + \
-//backends/xnnpack/... + //configurations/... + //kernels/portable/cpu/... + \
-//runtime/... + //schema/... + //test/... + //util/..."
+//backends/xnnpack/... + //configurations/... + //kernels/aten/... + \
+//kernels/optimized/... + //kernels/portable/... + //kernels/quantized/... + \
+//kernels/test/... + //runtime/... + //schema/... + //test/... + //util/..."
 
 # TODO: expand the covered scope of Buck targets.
-buck2 build //runtime/core/portable_type/...
-buck2 test //runtime/core/portable_type/...
+# //runtime/kernel/... is failing because //third-party:torchgen_files's shell script can't find python on PATH.
+# //runtime/test/... requires Python torch, which we don't have in our OSS buck setup.
+buck2 build //runtime/backend/... //runtime/core/... //runtime/executor: //runtime/kernel/... //runtime/platform/...
+buck2 test //runtime/backend/... //runtime/core/... //runtime/executor: //runtime/kernel/... //runtime/platform/...
diff --git a/runtime/core/exec_aten/testing_util/test/targets.bzl b/runtime/core/exec_aten/testing_util/test/targets.bzl
index 52a6dd0b7da..a37d08ecb22 100644
--- a/runtime/core/exec_aten/testing_util/test/targets.bzl
+++ b/runtime/core/exec_aten/testing_util/test/targets.bzl
@@ -1,4 +1,4 @@
-load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "runtime")
 
 def define_common_targets():
     """Defines targets that should be shared between fbcode and xplat.
@@ -15,19 +15,15 @@ def define_common_targets():
         ],
     )
 
-    runtime.cxx_test(
-        name = "tensor_factory_test",
-        srcs = ["tensor_factory_test.cpp"],
-        deps = [
-            "//executorch/runtime/core/exec_aten/testing_util:tensor_util",
-        ],
-    )
+    for aten_mode in get_aten_mode_options():
+        aten_suffix = "_aten" if aten_mode else ""
+        preprocessor_flags = ["-DUSE_ATEN_LIB"] if aten_mode else []
 
-    runtime.cxx_test(
-        name = "tensor_factory_test_aten",
-        srcs = ["tensor_factory_test.cpp"],
-        preprocessor_flags = ["-DUSE_ATEN_LIB"],
-        deps = [
-            "//executorch/runtime/core/exec_aten/testing_util:tensor_util_aten",
-        ],
-    )
+        runtime.cxx_test(
+            name = "tensor_factory_test" + aten_suffix,
+            srcs = ["tensor_factory_test.cpp"],
+            preprocessor_flags = preprocessor_flags,
+            deps = [
+                "//executorch/runtime/core/exec_aten/testing_util:tensor_util" + aten_suffix,
+            ],
+        )
diff --git a/runtime/core/test/targets.bzl b/runtime/core/test/targets.bzl
index ef09c905674..180e4eb0a0d 100644
--- a/runtime/core/test/targets.bzl
+++ b/runtime/core/test/targets.bzl
@@ -93,14 +93,15 @@ def define_common_targets():
         ],
     )
 
-    runtime.cxx_test(
-        name = "tensor_shape_dynamism_test_aten",
-        srcs = ["tensor_shape_dynamism_test_aten.cpp"],
-        deps = [
-            "//executorch/runtime/core/exec_aten:lib_aten",
-            "//executorch/runtime/core/exec_aten/testing_util:tensor_util_aten",
-        ],
-    )
+    if True in get_aten_mode_options():
+        runtime.cxx_test(
+            name = "tensor_shape_dynamism_test_aten",
+            srcs = ["tensor_shape_dynamism_test_aten.cpp"],
+            deps = [
+                "//executorch/runtime/core/exec_aten:lib_aten",
+                "//executorch/runtime/core/exec_aten/testing_util:tensor_util_aten",
+            ],
+        )
 
     for aten_mode in get_aten_mode_options():
         aten_suffix = "_aten" if aten_mode else ""