From dc0850541b67d05a4cc004c46b20ab98e494ee21 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Mon, 19 Apr 2021 14:56:24 -0500
Subject: [PATCH] Fix test update for mg_temperature and surface_source

---
 tests/regression_tests/mg_temperature/test.py | 59 ++-----------------
 tests/regression_tests/surface_source/test.py |  7 +++
 2 files changed, 13 insertions(+), 53 deletions(-)

diff --git a/tests/regression_tests/mg_temperature/test.py b/tests/regression_tests/mg_temperature/test.py
index 6bcc55492a..af66390b53 100644
--- a/tests/regression_tests/mg_temperature/test.py
+++ b/tests/regression_tests/mg_temperature/test.py
@@ -6,7 +6,7 @@ import shutil
 
 class MgTemperatureTestHarness(TestHarness):
 
-    def execute_test(self):
+    def execute_test(self, update=False):
         """Run OpenMC with the appropriate arguments and check the outputs."""
         base_dir = os.getcwd()
         overall_results = []
@@ -47,10 +47,12 @@ class MgTemperatureTestHarness(TestHarness):
                 overall_results.append(results)
             os.chdir(base_dir)
             self._write_results("".join(overall_results))
-            self._compare_results()
+            if update:
+                self._overwrite_results()
+            else:
+                self._compare_results()
         finally:
             os.chdir(base_dir)
-            shutil.copyfile("results_test.dat", "results_true.dat")
             if (os.path.isdir("./temp")):
                 shutil.rmtree("./temp")
             self._cleanup()
@@ -59,56 +61,7 @@ class MgTemperatureTestHarness(TestHarness):
                     os.remove(f)
 
     def update_results(self):
-        """Update the results_true using the current version of OpenMC."""
-        base_dir = os.getcwd()
-        overall_results = []
-        macro_xs = create_openmc_2mg_libs(names)
-        types = ('micro', 'micro',
-                 'micro', 'micro',
-                 'micro',
-                 'macro', 'macro',
-                 'macro', 'macro',
-                 'macro')
-        temperatures = (300., 600., 900.,
-                        520., 600.,
-                        300., 600., 900.,
-                        520., 600)
-        methods = 2 * (3 * ('nearest',) + 2 * ('interpolation',))
-        analyt_interp = 10 * [None]
-        analyt_interp[3] = (600. - 520.) / 300.
-        analyt_interp[8] = (600. - 520.) / 300.
-        try:
-            if (os.path.isdir("./temp")):
-                shutil.rmtree("./temp")
-            os.mkdir("temp")
-            os.chdir(os.path.join(base_dir, "temp"))
-            for cs, t, m, ai in zip(types, temperatures, methods, analyt_interp):
-                if (cs == 'macro'):
-                    build_inf_model(['macro'], '../macro_2g.h5', t, m)
-                else:
-                    build_inf_model(names, '../micro_2g.h5', t, m)
-                if not ai:
-                    kanalyt = analytical_solution_2g_therm(macro_xs[t])
-                else:
-                    kanalyt = analytical_solution_2g_therm(macro_xs[300],
-                                                           macro_xs[600], ai)
-                self._run_openmc()
-                self._test_output_created()
-                string = "{}, method: {}, t: {}, {}kanalyt\n{:12.6E}\n"
-                results = string.format(cs, m, t, self._get_results(), kanalyt)
-                overall_results.append(results)
-            os.chdir(base_dir)
-            self._write_results("".join(overall_results))
-            self._compare_results()
-        finally:
-            os.chdir(base_dir)
-            shutil.copyfile("results_test.dat", "results_true.dat")
-            if (os.path.isdir("./temp")):
-                shutil.rmtree("./temp")
-            self._cleanup()
-            for f in ['micro_2g.h5', 'macro_2g.h5']:
-                if os.path.exists(f):
-                    os.remove(f)
+        self.execute_test(update=True)
 
 
 def test_mg_temperature():
diff --git a/tests/regression_tests/surface_source/test.py b/tests/regression_tests/surface_source/test.py
index 492ef192e9..dc60cfd3c7 100644
--- a/tests/regression_tests/surface_source/test.py
+++ b/tests/regression_tests/surface_source/test.py
@@ -1,4 +1,5 @@
 import os
+import shutil
 
 import h5py
 import numpy as np
@@ -101,6 +102,12 @@ class SurfaceSourceTestHarness(PyAPITestHarness):
         finally:
             self._cleanup()
 
+    def _overwrite_results(self):
+        """Overwrite the results_true with the results_test."""
+        shutil.copyfile('results_test.dat', 'results_true.dat')
+        if os.path.exists('surface_source.h5'):
+            shutil.copyfile('surface_source.h5', 'surface_source_true.h5')
+
     def _cleanup(self):
         """Delete statepoints, tally, and test files."""
         super()._cleanup()