diff --git a/CMakeLists.txt b/CMakeLists.txt
index 57b93a56b35f80de29700d6e2c8ee241a83dd947..e3048e172ec5808808beffb81255608c3303d6f5 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -247,3 +247,12 @@ add_executable( inversion inversion.cpp ${EXTRA_INCS})
 target_link_libraries( inversion fwi  ${EXTRA_LIBS})
 add_executable( error error.cpp ${EXTRA_INCS})
 target_link_libraries( error fwi  ${EXTRA_LIBS})
+
+
+# Unit tests
+
+find_package(Boost 1.67.0 COMPONENTS unit_test_framework REQUIRED)
+include_directories(${Boost_INCLUDE_DIRS})
+
+add_executable(tests tests/test.cpp ${EXTRA_INCS})
+target_link_libraries(tests ${Boost_UNIT_TEST_FRAMEWORK_LIBRARY} ${EXTRA_LIBS} fwi)
diff --git a/common/wave/element.cpp b/common/wave/element.cpp
index 601b642186d18ecba28bb7f34d02c43362cf79d2..ab5974dfc68059fb64900f1a2d324649cb45760c 100644
--- a/common/wave/element.cpp
+++ b/common/wave/element.cpp
@@ -8,9 +8,9 @@ template<Physic T_Physic>
 void WaveMultiField<T_Physic>::initializeField(const WaveField<T_Physic>& field)
 {
     _field = field;
-    for (auto it = _dofValues.begin(); it != _dofValues.end(); it++)
+    for (unsigned i = 0; i < _dofValues.size(); ++i)
     {
-        it->resize(_field.numberOfDofs());
+      _dofValues[i].resize(_field.numberOfDofs());
     }
     _isInitialized = true;
 }
@@ -49,19 +49,29 @@ template<Physic T_Physic>
 void WaveAuxilaryField<T_Physic>::assignValues(const std::vector< std::complex<double> > &values)
 {
     if(_systemValuesToDofValues.size() == 0) {
-      _systemValuesToDofValues.reserve(values.size());
+      _systemValuesToDofValues.reserve(_wmf->_field.numberOfDofs());
+      _unknownIndices.reserve(values.size());
+      unsigned counter = 0;
       for(auto it = _wmf->_field.begin(); it != _wmf->_field.end(); it++)
       {
-          if(it->first->type() == gmshfem::dofs::Type::Unknown)
-          {
-              _systemValuesToDofValues.push_back(it->first->numDof() - 1);
-          }
+        // Save all DOFs, including the fixed ones, but remember to only update the unknown ones
+         _systemValuesToDofValues.push_back(it->first->numDof() - 1);
+        if (it->first->type() == gmshfem::dofs::Type::Unknown) {
+          _unknownIndices.push_back(counter);
+        }
+        ++counter;
+
       }
     }
-    for(unsigned int i = 0; i < values.size(); ++i) {
+
+    for (unsigned int i : _unknownIndices)
+    {
       _wmf->_dofValues[_index][i] = values[_systemValuesToDofValues[i]];
     }
-    if(_index == _wmf->_index){_wmf->_isUpToDate=false;}
+    if (_index == _wmf->_index)
+    {
+      _wmf->_isUpToDate = false;
+    }
 }
 
 template class WaveMultiField<Physic::acoustic>;
diff --git a/common/wave/element.h b/common/wave/element.h
index 5b1565c51d6fb031ddd456fc52b1a5ff74083fd3..7261eacbead3fa15b1f28617dd2047c00eff7d63 100644
--- a/common/wave/element.h
+++ b/common/wave/element.h
@@ -49,6 +49,7 @@ private:
     WaveMultiField<T_Physic>* _wmf;
     unsigned int _index;
     std::vector< unsigned int > _systemValuesToDofValues;
+    std::vector< unsigned int > _unknownIndices;
 public:
     WaveAuxilaryField(std::string name, gmshfem::domain::Domain domain, std::string gmodel, const wave::Discretization<T_Physic>& w_discret): WaveField<T_Physic>(name,domain,gmodel,w_discret), _wmf(nullptr), _index(0), _systemValuesToDofValues() {};
     void setField(WaveMultiField<T_Physic>* wmf){_wmf=wmf;_systemValuesToDofValues.clear();};
diff --git a/tests/test.cpp b/tests/test.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..90c0a351a8915bd26f239058458196d16bbf037b
--- /dev/null
+++ b/tests/test.cpp
@@ -0,0 +1,105 @@
+#define BOOST_TEST_MODULE fwi_tests
+#include <boost/test/unit_test.hpp>
+
+#include <cstring>
+
+#include <gmsh.h>
+#include <gmshfem/Formulation.h>
+#include <gmshfem/GmshFem.h>
+#include "../common/wave/element.h"
+
+#define INIT_GMSHFEM()                  \
+    namespace geo = gmsh::model::geo;   \
+    namespace model = gmsh::model;      \
+    namespace mesh = gmsh::model::mesh; \
+    using namespace gmshfem;            \
+    using namespace gmshfem::common;    \
+    using namespace gmshfem::problem;   \
+    using namespace gmshfem::domain;    \
+    using namespace gmshfem::field;     \
+    using namespace gmshfem::function;  \
+    using namespace gmshfem::post;      \
+    using namespace gmshfem::equation;  \
+    gmshfem::common::GmshFem fem(0, nullptr);
+
+BOOST_AUTO_TEST_SUITE(WaveMultiFieldTest);
+
+BOOST_AUTO_TEST_CASE(TestHomoegenuousDirichletWMF)
+{
+    INIT_GMSHFEM();
+    fem.getArgsManager()->addUserDefinedParameter("wave_FunctionSpaceDegree", "2");
+    fem.getArgsManager()->addUserDefinedParameter("wave_FunctionSpaceType", "HierarchicalH1");
+
+    double lc = 0.1;
+    model::add("Poisson");
+
+    geo::addPoint(0.0, 0.0, 0.0, lc, 1);
+    geo::addPoint(1.0, 0.0, 0.0, lc, 2);
+    geo::addPoint(1.0, 1.0, 0.0, lc, 3);
+    geo::addPoint(0.0, 1.0, 0.0, lc, 4);
+
+    geo::addLine(1, 2);
+    geo::addLine(2, 3);
+    geo::addLine(3, 4);
+    geo::addLine(4, 1);
+
+    geo::addCurveLoop({1,2,3,4});
+    geo::addPlaneSurface({1});
+    geo::synchronize();
+    mesh::generate();
+
+    model::addPhysicalGroup(1, {4}, 1);
+    model::setPhysicalName(1, 1, "gammaLeft");
+
+    model::addPhysicalGroup(2, {1}, 1);
+    model::setPhysicalName(2, 1, "omega");
+
+    Formulation<std::complex<double>> formulation("Poisson");
+    Domain omega("omega");
+    Domain gammaLeft("gammaLeft");
+    
+    // Allocate an Auxiliary Field that will serve as proxy to write in the WMF
+    WaveAuxilaryField<Physic::acoustic> v("v", omega, "Poisson", wave::Discretization<Physic::acoustic>(fem));
+    v.addConstraint(gammaLeft, 0.); // Apply Dirichlet *before* preprocessing, to get correct DOF dictionary
+
+    // Assemble LHS to get the correct number of DOFs
+    formulation.integral(grad(dof(v)), grad(tf(v)), omega, "Gauss6");
+    formulation.pre();
+
+    WaveMultiField<Physic::acoustic> wmf(2, "v_multi");
+    wmf.initializeField(v);
+    v.setField(&wmf);
+
+    // First solve
+    v.setIndex(0);
+    formulation.integral(-1., tf(v), omega, "Gauss6");
+    formulation.assemble();
+    formulation.solve(true);
+        
+    // Second solve
+    formulation.setRHSToZero();
+
+    v.setIndex(1);
+    formulation.removeTerms();
+    formulation.integral(-2., tf(v), omega, "Gauss6");
+    formulation.assemble();
+    formulation.solve(true);
+
+    // Check analytical solution to tolerance.
+    double f = 1;
+    auto expected = f * x<std::complex<double>>() * (1. - 0.5 * x<std::complex<double>>());
+    
+    // Quadratic solution -> should be exact on a coarse mesh.
+    auto integralV = abs(integrate(v, omega, "Gauss4"));
+    auto integralErr1 = abs(integrate(expected - wmf[0], omega, "Gauss4"));
+    auto integralErr2 = abs(integrate(2. * expected - wmf[1], omega, "Gauss4"));
+    BOOST_TEST(integralV < 1e-8, "AuxiliaryWaveField should be empty. Integral of absolute value is " << integralV);
+    BOOST_TEST(integralErr1 < 1e-8, "First WMF has the wrong value. Error norm: " << integralErr1);
+    BOOST_TEST(integralErr2 < 1e-8, "Second WMF has the wrong value. Error norm: " << integralErr2);
+}
+BOOST_AUTO_TEST_CASE(dummy_test)
+{
+    BOOST_CHECK(1 + 1 == 2);
+}
+
+BOOST_AUTO_TEST_SUITE_END()
\ No newline at end of file