diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
index 90f73a502..e77dff92e 100644
--- a/gusto/time_discretisation.py
+++ b/gusto/time_discretisation.py
@@ -216,14 +216,14 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = self.wrapper.subwrappers[field].function_space
                     
-                    # Replace with new test function
-                    # STILL TO SORT OUT.
-                    #new_test = TestFunction(self.wrapper.subwrappers[field].test_space)
-                    
-                    #self.residual = self.residual.label_map(
-                    #    all_terms,
-                    #    map_if_true=replace_test_function(new_test, old_idx=field_idx))
                     
+                    # Replace the test function space
+                    # This currently won't work: supg test 
+                    # is a function, not a space
+                    if self.wrapper.suboptions[field].name == "supg":
+                        self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test
+                    else:
+                        self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test_space
                 
                 self.wrapper.setup()
                 
@@ -244,7 +244,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                 
                 # Only call this if with SUPG, so should put this into the 
                 # previous section
-                self.residual = self.wrapper.label_terms(self.residual)
+                # self.residual = self.wrapper.label_terms(self.residual)
                 
                 
             else:
diff --git a/gusto/wrappers.py b/gusto/wrappers.py
index 0e6c655d5..052ddbc2e 100644
--- a/gusto/wrappers.py
+++ b/gusto/wrappers.py
@@ -423,6 +423,7 @@ def __init__(self, equation, suboptions):
         
         #self.wrapper_spaces = equation.space_names
         self.wrapper_spaces = equation.spaces
+        self.test_spaces = equation.spaces
         print(len(equation.spaces))
         
         #self.x_in = Function(equation.function_space)
@@ -475,8 +476,7 @@ def setup(self):
         self.function_space = MixedFunctionSpace(self.wrapper_spaces)
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space) 
-        self.test_space = self.function_space
-        #pass
+        self.test_space = MixedFunctionSpace(self.test_spaces)
     
     def pre_apply(self, x_in):
         """