review

openmdao/components/tests/test_balance_comp.py

@@ -897,26 +897,17 @@ def test_specified_shape(self):
         assert_check_partials(cpd, atol=1e-5, rtol=1e-5)
 
     def test_shape_from_rhs_val(self):
-        n = 5
         p = om.Problem()
-
-        p.model.add_subsystem('comp_A',
-                              om.ExecComp('y = x**2', y={'shape': (n, )}, x={'shape': (n, )}))
-        p.model.add_subsystem('comp_B',
-                              om.ExecComp('y = x + 7', y={'shape': (n, )}, x={'shape': (n, )}))
-
-        p.model.add_subsystem('bal', om.BalanceComp('x', rhs_val=np.ones((n, ))))
-        p.model.connect('bal.x', 'comp_A.x')
-        p.model.connect('bal.x', 'comp_B.x')
-        p.model.connect('comp_A.y', 'bal.lhs:x')
-        p.model.connect('comp_B.y', 'bal.rhs:x')
-
-        p.model.linear_solver = om.DirectSolver()
-        p.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
+        init = np.ones((5, ))
+        p.model.add_subsystem('bal', om.BalanceComp('x', rhs_val=init))
 
         # Bug was a size mismatch exception raised during setup.
         p.setup()
 
+        self.assertTrue(p.get_val('bal.x').shape == init.shape)
+        self.assertTrue(p.get_val('bal.lhs:x').shape == init.shape)
+        self.assertTrue(p.get_val('bal.rhs:x').shape == init.shape)
+
 
 if __name__ == '__main__':  # pragma: no cover
     unittest.main()