diff --git a/unit_tests/operator/Test_ExpandImpl.cpp b/unit_tests/operator/Test_ExpandImpl.cpp
index 878c608110eabb824d8a6c0d1ceb0853b3c1449d..ad30457d33307ca595ecddfd3b06d58e118a02d0 100644
--- a/unit_tests/operator/Test_ExpandImpl.cpp
+++ b/unit_tests/operator/Test_ExpandImpl.cpp
@@ -13,20 +13,20 @@
 
 #include <catch2/catch_test_macros.hpp>
 
-#include "aidge/backend/cpu/data/TensorImpl.hpp"
-#include "aidge/backend/cpu/operator/ExpandImpl.hpp"
 #include "aidge/data/DataType.hpp"
 #include "aidge/data/Tensor.hpp"
 #include "aidge/operator/Expand.hpp"
 #include "aidge/utils/ArrayHelpers.hpp"
 
-using std::shared_ptr;
 
-using namespace Aidge;
+namespace Aidge {
+
+using std::shared_ptr;
 
-void setupTestExpand(shared_ptr<Tensor> inputData,
-                     shared_ptr<Tensor> inputShape,
-                     shared_ptr<Expand_Op> &op) {
+static void setupTestExpand(shared_ptr<Tensor> inputData,
+                            shared_ptr<Tensor> inputShape,
+                            shared_ptr<Expand_Op> &op,
+                            Tensor &expectedOutput) {
 
     op->getOutput(0)->setDataType(inputData->dataType());
 
@@ -35,6 +35,9 @@ void setupTestExpand(shared_ptr<Tensor> inputData,
 
     inputShape->setBackend("cpu");
     op->associateInput(1, inputShape);
+
+    expectedOutput.setBackend("cpu");
+    expectedOutput.setDataType(DataType::Int32);
 }
 
 TEST_CASE("[cpu/operator] Expand(forward)", "[Expand][CPU]") {
@@ -49,7 +52,7 @@ TEST_CASE("[cpu/operator] Expand(forward)", "[Expand][CPU]") {
             Array4D<cpptype_t<DataType::Int32>, 1, 3, 4, 2>({{{{{1, 3}, {1, 3}, {1, 3}, {1, 3}},
                                         {{1, 3}, {1, 3}, {1, 3}, {1, 3}},
                                         {{1, 3}, {1, 3}, {1, 3}, {1, 3}}}}});
-        setupTestExpand(inputData, inputShape, op);
+        setupTestExpand(inputData, inputShape, op, expectedOutput);
 
         // forwardDims has already been tested in core
         CHECK(op->forwardDims(true));
@@ -63,7 +66,7 @@ TEST_CASE("[cpu/operator] Expand(forward)", "[Expand][CPU]") {
             std::make_shared<Tensor>(Array1D<std::int64_t, 2>({2, 3}));
         Tensor expectedOutput = Array3D<cpptype_t<DataType::Int32>, 2, 2, 3>(
             {{{{2, 1, 3}, {2, 1, 3}}, {{2, 1, 3}, {2, 1, 3}}}});
-        setupTestExpand(inputData, inputShape, op);
+        setupTestExpand(inputData, inputShape, op,expectedOutput);
 
         // forwardDims has already been tested in core
         CHECK(op->forwardDims(true));
@@ -77,7 +80,7 @@ TEST_CASE("[cpu/operator] Expand(forward)", "[Expand][CPU]") {
             std::make_shared<Tensor>(Array1D<std::int64_t, 1>({1}));
         Tensor expectedOutput =
             Array4D<cpptype_t<DataType::Int32>, 2, 1, 3, 1>({{{2, 1, 3}, {2, 1, 3}}});
-        setupTestExpand(inputData, inputShape, op);
+        setupTestExpand(inputData, inputShape, op, expectedOutput);
 
         // forwardDims has already been tested in core
         CHECK(op->forwardDims(true));
@@ -91,7 +94,7 @@ TEST_CASE("[cpu/operator] Expand(forward)", "[Expand][CPU]") {
             std::make_shared<Tensor>(Array1D<std::int64_t, 3>({2, 1, 1}));
         Tensor expectedOutput =
             Array4D<cpptype_t<DataType::Int32>, 1, 2, 3, 1>({{{{2, 1, 3}, {2, 1, 3}}}});
-        setupTestExpand(inputData, inputShape, op);
+        setupTestExpand(inputData, inputShape, op,expectedOutput);
 
         // forwardDims has already been tested in core
         CHECK(op->forwardDims(true));
@@ -101,3 +104,4 @@ TEST_CASE("[cpu/operator] Expand(forward)", "[Expand][CPU]") {
     SECTION("N-Dim to N-Dim") {}
     auto inputData = std::shared_ptr<Tensor>();
 }
+} // namespace Aidge