diff --git a/dymos/examples/finite_burn_orbit_raise/test/test_ex_two_burn_orbit_raise.py b/dymos/examples/finite_burn_orbit_raise/test/test_ex_two_burn_orbit_raise.py
index a0836a0bf..9eb2dc2a7 100644
--- a/dymos/examples/finite_burn_orbit_raise/test/test_ex_two_burn_orbit_raise.py
+++ b/dymos/examples/finite_burn_orbit_raise/test/test_ex_two_burn_orbit_raise.py
@@ -263,9 +263,9 @@ def two_burn_orbit_raise_problem(transcription='gauss-lobatto', optimizer='SLSQP
     return p
 
 
+@use_tempdirs
 class TestExampleTwoBurnOrbitRaise(unittest.TestCase):
 
-    @use_tempdirs
     def test_ex_two_burn_orbit_raise(self):
         _, optimizer = set_pyoptsparse_opt('SNOPT', fallback=False)
 
@@ -279,9 +279,9 @@ def test_ex_two_burn_orbit_raise(self):
 
 
 # This test is separate because connected phases aren't directly parallelizable.
+@use_tempdirs
 class TestExampleTwoBurnOrbitRaiseConnected(unittest.TestCase):
 
-    @use_tempdirs
     def test_ex_two_burn_orbit_raise_connected(self):
         _, optimizer = set_pyoptsparse_opt('SNOPT', fallback=False)
 
diff --git a/dymos/trajectory/trajectory.py b/dymos/trajectory/trajectory.py
index 78f1dcc06..12cd33c21 100644
--- a/dymos/trajectory/trajectory.py
+++ b/dymos/trajectory/trajectory.py
@@ -372,12 +372,21 @@ def _setup_linkages(self):
                         p2.set_time_options(input_initial=True)
                 else:
                     vars_to_constrain.append(var)
-                    if var in p1_states:
+                    if var in p1.state_options:
                         units_map[var] = p1.state_options[var]['units']
                         shape_map[var] = p1.state_options[var]['shape']
-                    elif var in p1_controls:
+                    elif var in p1.control_options:
                         units_map[var] = p1.control_options[var]['units']
                         shape_map[var] = p1.control_options[var]['shape']
+                    elif var in p1.polynomial_control_options:
+                        units_map[var] = p1.polynomial_control_options[var]['units']
+                        shape_map[var] = p1.polynomial_control_options[var]['shape']
+                    elif var in p1.design_parameter_options:
+                        units_map[var] = p1.design_parameter_options[var]['units']
+                        shape_map[var] = p1.design_parameter_options[var]['shape']
+                    elif var in p1.input_parameter_options:
+                        units_map[var] = p1.input_parameter_options[var]['units']
+                        shape_map[var] = p1.input_parameter_options[var]['shape']
                     elif var == 'time':
                         units_map[var] = p1.time_options['units']
                         shape_map[var] = (1,)