From 32bea6890b3fa60e5fe933144361850c3c55e05e Mon Sep 17 00:00:00 2001 From: Julien Barreau Date: Thu, 1 Feb 2024 14:48:07 +0100 Subject: [PATCH] fix(csv): export fails when there is more fields in data than projection --- src/agent_toolkit/forestadmin/agent_toolkit/utils/csv.py | 2 +- src/agent_toolkit/tests/utils/test_csv.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/src/agent_toolkit/forestadmin/agent_toolkit/utils/csv.py b/src/agent_toolkit/forestadmin/agent_toolkit/utils/csv.py index d29abfed3..125316da2 100644 --- a/src/agent_toolkit/forestadmin/agent_toolkit/utils/csv.py +++ b/src/agent_toolkit/forestadmin/agent_toolkit/utils/csv.py @@ -16,7 +16,7 @@ class Csv: def make_csv(rows: List[Dict[str, Any]], projection: Projection) -> str: try: dumped_csv = StringIO() - csv_writer = DictWriter(dumped_csv, fieldnames=projection) + csv_writer = DictWriter(dumped_csv, fieldnames=projection, extrasaction="ignore") csv_writer.writeheader() for row in rows: csv_writer.writerow(Csv.format_field(row)) diff --git a/src/agent_toolkit/tests/utils/test_csv.py b/src/agent_toolkit/tests/utils/test_csv.py index 901dbff06..7e7a21813 100644 --- a/src/agent_toolkit/tests/utils/test_csv.py +++ b/src/agent_toolkit/tests/utils/test_csv.py @@ -62,3 +62,10 @@ def test_make_csv_error(self): self.assertRaises(CsvException, Csv.make_csv, copy.deepcopy(self.data), self.data[0].keys()) mocked_write_row.assert_called_once() + + def test_csv_should_have_only_projection_fields_even_if_more_in_record(self): + fieldnames = [*self.data[0].keys()] + fieldnames.remove("boolean_field") + ret = Csv.make_csv(copy.deepcopy(self.data), fieldnames) + data = self.read_csv(ret) + self.assertNotIn("boolean_field", data[0].keys())