From 9d30ad51d0becbe45018f74b80867ecd0b9dad58 Mon Sep 17 00:00:00 2001
From: Alasdair Gray <alachris@umich.edu>
Date: Sat, 10 Feb 2024 11:49:34 -0500
Subject: [PATCH] Partial fix for check_totals issue (CS passes in serial now)

---
 .../openmdao_analysis_base_test.py            | 30 +++++++++----------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/tests/integration_tests/openmdao_analysis_base_test.py b/tests/integration_tests/openmdao_analysis_base_test.py
index d74d585bf..d9b20ca8a 100644
--- a/tests/integration_tests/openmdao_analysis_base_test.py
+++ b/tests/integration_tests/openmdao_analysis_base_test.py
@@ -14,16 +14,16 @@
 
 """
 This is a base class for running openmdao unit test cases.
-This base class will test function evaluations, partial, and total 
+This base class will test function evaluations, partial, and total
 sensitivities for the user-specified openmdao problems implemented by
-the child test case. When the user creates a new test based 
-on this class two methods are required to be defined in the child class. 
+the child test case. When the user creates a new test based
+on this class two methods are required to be defined in the child class.
     1. setup_problem
     2. setup_funcs
 
-See the virtual method implementations for each method 
+See the virtual method implementations for each method
 below for more details.
-NOTE: The child class must NOT implement its own setUp method 
+NOTE: The child class must NOT implement its own setUp method
 for the unittest class. This is handled in the base class.
 """
 
@@ -116,25 +116,25 @@ def test_totals(self):
             """
             Test total sensitivities using fd/cs
             """
+            of = list(self.func_ref.keys())
             # solve
             self.prob.run_model()
 
-            # Test functions total sensitivities
-            of = list(self.func_ref.keys())
-            for var_wrt in self.wrt:
-                with self.subTest(wrt=var_wrt):
-                    for var_of in of:
-                        with self.subTest(of=var_of):
-                            data = self.prob.check_totals(
-                                of=var_of,
-                                wrt=var_wrt,
+            data = self.prob.check_totals(
+                                of=of,
+                                wrt=self.wrt,
                                 compact_print=True,
                                 out_stream=None,
                                 method=self.fd_method,
                                 form=self.fd_form,
                                 step=self.dh,
                             )
-                            assert_check_totals(data, atol=self.atol, rtol=self.rtol)
+
+            # Test functions total sensitivities
+            for var_of, var_wrt in data.keys():
+                with self.subTest(of=var_of, wrt=var_wrt):
+                    test_data = {(var_of, var_wrt): data[(var_of, var_wrt)]}
+                    assert_check_totals(test_data, atol=self.atol, rtol=self.rtol)
 
         def cleanup_fwd_data(self, data):
             """