Updated testing for csv writer (#268)

SABS-R3-Epidemiology · Aug 6, 2024 · 222cbfb · 222cbfb
1 parent 8b33fa1
commit 222cbfb
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 11 deletions.
diff --git a/pyEpiabm/pyEpiabm/tests/test_unit/test_output/test_csv_dict_writer.py b/pyEpiabm/pyEpiabm/tests/test_unit/test_output/test_csv_dict_writer.py
@@ -20,7 +20,7 @@ def test_init(self, mock_mkdir):
                                          mock_categories)
             del m
         mo.assert_called_once_with(
-            os.path.join('mock_folder', 'mock_filename'), 'w')
+            os.path.join('mock_folder', 'mock_filename'), 'w', newline='')
         mo().write.assert_called_once_with('Cat1,Cat2,Cat3\r\n')
         mock_mkdir.assert_called_with('mock_folder')
 

diff --git a/pyEpiabm/pyEpiabm/tests/test_unit/test_routine/test_simulation.py b/pyEpiabm/pyEpiabm/tests/test_unit/test_routine/test_simulation.py
@@ -83,7 +83,7 @@ def test_configure(self, mock_mkdir):
             del test_sim.ih_status_writer
             del test_sim.ih_infectiousness_writer
             del test_sim.secondary_infections_writer
-        mo.assert_called_with(filename, 'w')
+        mo.assert_called_with(filename, 'w', newline='')
 
     @patch('os.makedirs')
     @patch('logging.warning')
@@ -132,7 +132,7 @@ def test_configure_ih_status(self, mock_mkdir):
             del test_sim.ih_status_writer
             del test_sim.ih_infectiousness_writer
             del test_sim.secondary_infections_writer
-        mo.assert_called_with(filename, 'w')
+        mo.assert_called_with(filename, 'w', newline='')
 
     @patch('os.makedirs')
     def test_configure_ih_infectiousness(self, mock_mkdir):
@@ -161,7 +161,7 @@ def test_configure_ih_infectiousness(self, mock_mkdir):
             del test_sim.ih_status_writer
             del test_sim.ih_infectiousness_writer
             del test_sim.secondary_infections_writer
-        mo.assert_called_with(filename, 'w')
+        mo.assert_called_with(filename, 'w', newline='')
 
     @patch('os.makedirs')
     def test_configure_secondary_infections(self, mock_mkdir):
@@ -191,7 +191,7 @@ def test_configure_secondary_infections(self, mock_mkdir):
             del test_sim.ih_status_writer
             del test_sim.ih_infectiousness_writer
             del test_sim.secondary_infections_writer
-        mo.assert_called_with(filename, 'w')
+        mo.assert_called_with(filename, 'w', newline='')
 
     @patch('logging.exception')
     @patch('os.path.join')

diff --git a/pyEpiabm/pyEpiabm/tests/test_unit/test_sweep/test_initial_demographics_sweep.py b/pyEpiabm/pyEpiabm/tests/test_unit/test_sweep/test_initial_demographics_sweep.py
@@ -1,4 +1,5 @@
 import os
+import unittest
 from unittest.mock import patch, mock_open, call
 
 import pyEpiabm as pe
@@ -87,8 +88,8 @@ def test_construct(self, mock_mkdir):
                                                       "Age group 0"])
             del dem_sweep.writer
             del dem_sweep.counts_writer
-        mo.assert_has_calls([call(file_name, 'w'),
-                             call(counts_file_name, 'w')],
+        mo.assert_has_calls([call(file_name, 'w', newline=''),
+                             call(counts_file_name, 'w', newline='')],
                             any_order=True)
 
     @patch('os.makedirs')
@@ -115,8 +116,8 @@ def test_construct_age_spatial(self, mock_mkdir):
                              [f"Age group {i}" for i in range(17)])
             del dem_sweep.writer
             del dem_sweep.counts_writer
-        mo.assert_has_calls([call(file_name, 'w'),
-                             call(counts_file_name, 'w')],
+        mo.assert_has_calls([call(file_name, 'w', newline=''),
+                             call(counts_file_name, 'w', newline='')],
                             any_order=True)
 
     @patch('os.makedirs')
@@ -253,3 +254,7 @@ def test_write_to_file_age_spatial(self, mock_mkdir):
         mock_mkdir.assert_called_with(os.path.join(os.getcwd(),
                                                    self.dem_file_params[
                                                        "output_dir"]))
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/python_examples/spatial_example/spatial_simulation_flow.py b/python_examples/spatial_example/spatial_simulation_flow.py
@@ -33,7 +33,8 @@
 
 # Version I: Create a population framework and save to file.
 population = pe.routine.ToyPopulationFactory.make_pop(pop_params)
-pe.routine.ToyPopulationFactory.assign_cell_locations(population, method='grid')
+pe.routine.ToyPopulationFactory.assign_cell_locations(population,
+                                                      method='grid')
 pe.routine.FilePopulationFactory.print_population(population, file_loc)
 
 # Version II: Generate population from input file.
@@ -84,8 +85,9 @@
 
 df = df.pivot(index="time", columns="cell",
               values="InfectionStatus.InfectMild")
-df.plot(); plt.legend().remove()
+df.plot()
 
+plt.legend().remove()
 plt.title("Infection curves for multiple cells")
 plt.ylabel("Infected Population")
 plt.savefig(