From cbf1f65a8f3389f55bb55249ce8c358ac06cf75c Mon Sep 17 00:00:00 2001
From: Joshua Larsen <jlarsen@usgs.gov>
Date: Thu, 14 May 2020 06:50:43 -0700
Subject: [PATCH] Fix(#870): Update ims name file record after removing ims
 package (#880)

* Fix(#870): Update ims name file record after removing ims package

* remove cbc test case

* update t504_test.py

* codacy update

* Update test_replace_ims_package for travis
---
 autotest/t504_test.py             | 36 +++++++++++++++++++++++--------
 flopy/mf6/modflow/mfsimulation.py | 26 ++++++++++++++++++++++
 2 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/autotest/t504_test.py b/autotest/t504_test.py
index 1f46b64358..c7c485037c 100644
--- a/autotest/t504_test.py
+++ b/autotest/t504_test.py
@@ -921,14 +921,32 @@ def test027_timeseriestest():
         assert pymake.compare_heads(None, None, files1=head_file, files2=head_new)
 
 
+def test_replace_ims_package():
+    pth = os.path.join(cpth, "test001e_UZF_3lay")
+    sim = flopy.mf6.MFSimulation.load("mfsim", sim_ws=pth, exe_name=exe_name)
+
+    ims = sim.ims
+    sim.remove_package(ims)
+
+    ims = flopy.mf6.ModflowIms(sim, print_option='SUMMARY',
+                               complexity="COMPLEX")
+    sim.register_ims_package(ims, ["gwf_1", ])
+    sim.write_simulation()
+    success, buff = sim.run_simulation()
+
+    if not success:
+        raise AssertionError()
+
+
 if __name__ == '__main__':
-    test001a_tharmonic()
+    # test001a_tharmonic()
     test001e_uzf_3lay()
-    test003_gwfs_disv()
-    test005_advgw_tidal()
-    test006_2models_mvr()
-    test006_gwf3()
-    test027_timeseriestest()
-    test036_twrihfb()
-    test045_lake1ss_table()
-    test045_lake2tr()
+    # test003_gwfs_disv()
+    # test005_advgw_tidal()
+    # test006_2models_mvr()
+    # test006_gwf3()
+    # test027_timeseriestest()
+    # test036_twrihfb()
+    # test045_lake1ss_table()
+    # test045_lake2tr()
+    test_replace_ims_package()
\ No newline at end of file
diff --git a/flopy/mf6/modflow/mfsimulation.py b/flopy/mf6/modflow/mfsimulation.py
index 11889720a4..22e0d50001 100644
--- a/flopy/mf6/modflow/mfsimulation.py
+++ b/flopy/mf6/modflow/mfsimulation.py
@@ -1313,6 +1313,7 @@ def remove_package(self, package_name):
                 del self._exchange_files[package.filename]
             if package.filename in self._ims_files:
                 del self._ims_files[package.filename]
+                self._remove_ims_soultion_group(package.filename)
             if package.filename in self._ghost_node_files:
                 del self._ghost_node_files[package.filename]
             if package.filename in self._mover_files:
@@ -1748,6 +1749,31 @@ def _get_package_path(package):
         else:
             return (package.package_type,)
 
+    def _remove_ims_soultion_group(self, ims_file):
+        solution_recarray = self.name_file.solutiongroup
+        for solution_group_num in solution_recarray.get_active_key_list():
+            try:
+                rec_array = solution_recarray.get_data(solution_group_num[0])
+            except MFDataException as mfde:
+                message = 'An error occurred while getting solution group' \
+                          '"{}" from the simulation name file' \
+                          '.'.format(solution_group_num[0])
+                raise MFDataException(mfdata_except=mfde,
+                                      package='nam',
+                                      message=message)
+
+            new_array = []
+            for record in rec_array:
+                if record.slnfname == ims_file:
+                    continue
+                else:
+                    new_array.append(record)
+
+            if not new_array:
+                new_array = None
+
+            solution_recarray.set_data(new_array, solution_group_num[0])
+
     def _append_to_ims_solution_group(self, ims_file, new_models):
         solution_recarray = self.name_file.solutiongroup
         for solution_group_num in solution_recarray.get_active_key_list():