#2670 fix unit tests

pybamm-team · Apr 29, 2023 · 4caac61 · 4caac61
1 parent ce514b1
commit 4caac61
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 14 deletions.
diff --git a/tests/unit/test_models/test_submodels/test_effective_current_collector.py b/tests/unit/test_models/test_submodels/test_effective_current_collector.py
@@ -75,15 +75,18 @@ def test_get_processed_variables(self):
         # for each current collector model
         for model in models[1:]:
             solution = model.default_solver.solve(model)
-            vars = model.post_process(solution, param, V, I)
+            variables = model.post_process(solution, param, V, I)
             pts = np.array([0.1, 0.5, 0.9]) * min(
                 param.evaluate(model.param.L_y), param.evaluate(model.param.L_z)
             )
-            for var, processed_var in vars.items():
+            for var, processed_var in variables.items():
                 if "Voltage [V]" in var:
                     processed_var(t=solution_1D.t[5])
                 else:
-                    processed_var(t=solution_1D.t[5], y=pts, z=pts)
+                    if model.options["dimensionality"] == 1:
+                        processed_var(t=solution_1D.t[5], z=pts)
+                    else:
+                        processed_var(t=solution_1D.t[5], y=pts, z=pts)
 
 
 if __name__ == "__main__":

diff --git a/tests/unit/test_plotting/test_quick_plot.py b/tests/unit/test_plotting/test_quick_plot.py
@@ -329,19 +329,15 @@ def test_loqs_spme(self):
                 )
                 quick_plot.plot(0)
 
-                qp_data = (
-                    quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][
-                        0
-                    ].get_ydata(),
-                )[0]
+                qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][
+                    0
+                ].get_ydata()
                 np.testing.assert_array_almost_equal(qp_data, c_e[:, 0])
-                quick_plot.slider_update(t_eval[-1] / scale)
 
-                qp_data = (
-                    quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][
-                        0
-                    ].get_ydata(),
-                )[0][:, 0]
+                quick_plot.slider_update(t_eval[-1] / scale)
+                qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][
+                    0
+                ].get_ydata()
                 np.testing.assert_array_almost_equal(qp_data, c_e[:, 1])
 
             # test quick plot of particle for spme