diff --git a/src/onegov/election_day/utils/archive_generator.py b/src/onegov/election_day/utils/archive_generator.py index b6f0677f69..a949bd1459 100644 --- a/src/onegov/election_day/utils/archive_generator.py +++ b/src/onegov/election_day/utils/archive_generator.py @@ -10,8 +10,9 @@ from onegov.core.csv import convert_list_of_dicts_to_csv from onegov.core.utils import module_path -from onegov.ballot import Vote, Election, ElectionCompound +from onegov.ballot import Vote, Election, ElectionCompound, ProporzElection from onegov.election_day.formats import export_internal +from onegov.election_day.formats import export_parties_internal from typing import Any @@ -82,12 +83,7 @@ def generate_csv(self) -> None: year_dir = f'{entity_name}/{year}' self.temp_fs.makedirs(year_dir, recreate=True) for item in yearly_package: - # item may be of type Vote, Election or ElectionCompound - filename = item.id[: self.MAX_FILENAME_LENGTH] + '.csv' - combined_path = path.combine(year_dir, filename) - with self.temp_fs.open(combined_path, 'w') as f: - rows = export_internal(item, sorted(self.app.locales)) - f.write(convert_list_of_dicts_to_csv(rows)) + self.export_item(item, year_dir) # Additionally, create 'flat csv' containing all votes in a single file if votes: @@ -223,6 +219,31 @@ def include_docs(self) -> None: dst_path=match.path, ) + def export_item(self, item: 'EntityT', dir: str) -> None: + locales = sorted(self.app.locales) + default_locale = self.app.default_locale + + # results + filename = item.id[:self.MAX_FILENAME_LENGTH] + '.csv' + combined_path = path.combine(dir, filename) + rows = export_internal(item, locales) + with self.temp_fs.open(combined_path, 'w') as f: + f.write(convert_list_of_dicts_to_csv(rows)) + + # party results + if getattr(item, 'has_party_results', False): + assert isinstance(item, (ProporzElection, ElectionCompound)) + filename = item.id[:self.MAX_FILENAME_LENGTH + 8] + '-parties.csv' + combined_path = path.combine(dir, filename) + rows = export_parties_internal( + item, + locales, + # FIXME: Should we assert that the default_locale is set? + default_locale=default_locale, # type:ignore[arg-type] + ) + with self.temp_fs.open(combined_path, 'w') as f: + f.write(convert_list_of_dicts_to_csv(rows)) + def generate_archive(self) -> str | None: self.generate_csv() self.include_docs() diff --git a/tests/onegov/election_day/utils/test_archive_generator.py b/tests/onegov/election_day/utils/test_archive_generator.py index ddcf0238d5..7d5e89a2f2 100644 --- a/tests/onegov/election_day/utils/test_archive_generator.py +++ b/tests/onegov/election_day/utils/test_archive_generator.py @@ -224,7 +224,7 @@ def test_long_filenames_are_truncated(election_day_app_zg): def test_election_generation(election_day_app_zg, import_test_datasets): - import_test_datasets( + election, errors = import_test_datasets( 'internal', 'election', 'zg', @@ -235,6 +235,17 @@ def test_election_generation(election_day_app_zg, import_test_datasets): dataset_name='nationalratswahlen-2015', app_session=election_day_app_zg.session() ) + assert not errors + errors = import_test_datasets( + 'internal', + 'parties', + 'zg', + 'canton', + 'proporz', + election=election, + dataset_name='nationalratswahlen-2015-parteien', + ) + assert not errors archive_generator = ArchiveGenerator(election_day_app_zg) zip_path = archive_generator.generate_archive() @@ -244,13 +255,14 @@ def test_election_generation(election_day_app_zg, import_test_datasets): assert "elections" in top_level_dir elections = zip_fs.opendir("elections") - elections_by_year = [year for year in elections.listdir(".")] - - assert {"2015"} == set(elections_by_year) - - csv = [csv for csv in zip_fs.scandir("elections/2015", - namespaces=["basic"])] - first_file = csv[0] - - assert "proporz_internal_nationalratswahlen-2015.csv" ==\ - first_file.name + years = {year for year in elections.listdir(".")} + assert years == {"2015"} + + files = { + csv.name for csv + in zip_fs.scandir("elections/2015", namespaces=["basic"]) + } + assert files == { + 'proporz_internal_nationalratswahlen-2015.csv', + 'proporz_internal_nationalratswahlen-2015-parties.csv' + }