From 756dfd55916b90f2bba759b9896bfbaa76ff7ee6 Mon Sep 17 00:00:00 2001
From: why-not-try-calmer <mrnycticorax@gmail.com>
Date: Wed, 31 Jan 2024 17:31:59 +0100
Subject: [PATCH] Importing...

---
 comptages/core/importer.py                          |  7 ++++---
 comptages/core/report.py                            |  2 +-
 .../datamodel/management/commands/importdata.py     |  8 ++++++--
 comptages/test/test_report.py                       | 13 +++++++++----
 comptages/test/test_statistics.py                   |  2 +-
 5 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/comptages/core/importer.py b/comptages/core/importer.py
index f73ddcc..51b15e5 100644
--- a/comptages/core/importer.py
+++ b/comptages/core/importer.py
@@ -9,8 +9,9 @@
 from comptages.core.bulk_create_manager import BulkCreateManager
 
 
-def simple_print_callback(progress):
-    print(f"Importing... {progress}%")
+def simple_print_callback(progress: int):
+    if progress % 10 == 0:
+        print(f"Importing... {progress}%")
 
 
 def import_file(file_path: str, count, callback_progress=simple_print_callback):
@@ -337,7 +338,7 @@ def _parse_file_header(file_path: str):
                 elif file_header["CLASS"][:5] == "FHWA ":
                     file_header["CLASS"] = "FHWA13"
                 elif file_header["CLASS"] == "CAT-Cycle_dist-empat":
-                    file_header["CLASS"] = "SPCH-MD_5C"
+                    file_header["CLASS"] = "SPCH-MD 5C"
 
     return file_header
 
diff --git a/comptages/core/report.py b/comptages/core/report.py
index e58a648..1f7388f 100644
--- a/comptages/core/report.py
+++ b/comptages/core/report.py
@@ -919,7 +919,7 @@ def _remove_useless_sheets(count: models.Count, workbook: Workbook):
 
     for key in to_remove_from_spreadsheet:
         try:
-            workbook.remove_sheet(workbook[key])
+            workbook.remove(workbook[key])
         except KeyError:
             continue
 
diff --git a/comptages/datamodel/management/commands/importdata.py b/comptages/datamodel/management/commands/importdata.py
index ac5a378..7ee9470 100644
--- a/comptages/datamodel/management/commands/importdata.py
+++ b/comptages/datamodel/management/commands/importdata.py
@@ -415,7 +415,7 @@ def import_municipalities(self, csv_file):
         Municipality.objects.bulk_create(objects)
         print(f"Inserted {len(objects)} municipalities.")
 
-    def import_count(self, limit: Optional[int] = None):
+    def import_count(self, limit=50):
         section_id = "00107695"
         year = 2021
         installations = Installation.objects.filter(lane__id_section=section_id)
@@ -444,7 +444,11 @@ def import_count(self, limit: Optional[int] = None):
         files = list(path_to_files.iterdir())
 
         path_to_files = Path("/OpenComptage/test_data/SWISS10_vbv_year")
-        files = list(path_to_files.iterdir())[:50]
+
+        if limit == 0:
+            files = list(path_to_files.iterdir())
+        else:
+            files = list(path_to_files.iterdir())[:limit]
 
         imported = 0
         for file in files:
diff --git a/comptages/test/test_report.py b/comptages/test/test_report.py
index 2c06780..4eea5cb 100644
--- a/comptages/test/test_report.py
+++ b/comptages/test/test_report.py
@@ -72,10 +72,15 @@ def test_all_sections_default(self):
             id_installation=installation,
         )
 
-        importer.import_file(test_utils.test_data_path("00056520.V01"), count)
-        importer.import_file(test_utils.test_data_path("00056520.V02"), count)
-
-        report.prepare_reports("/tmp/", count)
+        for file in Path(test_utils.test_data_path(test_data_folder)).iterdir():
+            importer.import_file(test_utils.test_data_path(str(file)), count)
+
+        report.prepare_reports(self.test_outputs, count)
+        found_files = len(list(Path(self.test_outputs).iterdir()))
+        # The number of files generated is expected to be: weeks measured x sections
+        # so let's make sure all sections are considered in the files generation
+        self.assertGreater(found_files, 0)
+        self.assertEqual(found_files % n_sections, 0)
 
     def test_busiest_by_season(self):
         # Import test data pertaining to "mobilité douce"
diff --git a/comptages/test/test_statistics.py b/comptages/test/test_statistics.py
index de7ce00..d041ddc 100644
--- a/comptages/test/test_statistics.py
+++ b/comptages/test/test_statistics.py
@@ -272,7 +272,7 @@ def test_get_speed_data_empty(self):
         self.assertFalse(statistics.get_speed_data(count, sections[2]).empty)
 
     def test_get_valid_days(self):
-        call_command("importdata", "--only-count")
+        call_command("importdata", "--only-count", "--limit=0")
         section_id = "00107695"
         section = models.Section.objects.get(id=section_id)
         # 2021 comes from the call to `importdata --only count`