From 1608d01bb68c1f6292b04c70caa609d34943b371 Mon Sep 17 00:00:00 2001 From: Alex Schofield Date: Wed, 19 Feb 2025 15:37:57 +0000 Subject: rename write function & update references --- test/test_write.py | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 test/test_write.py (limited to 'test/test_write.py') diff --git a/test/test_write.py b/test/test_write.py new file mode 100644 index 0000000..eceac28 --- /dev/null +++ b/test/test_write.py @@ -0,0 +1,57 @@ +import io +import csv +from obfuscator.csv_writer import create_byte_stream + + +def csv_bytes_to_list(csv_bytes: bytes): + csv_string = csv_bytes.decode("utf-8") + f = io.StringIO(csv_string) + reader = csv.DictReader(f) + return [dict(row) for row in reader] + + +def test_create_byte_stream_valid_data(): + data = [ + {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, + {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, + ] + csv_bytes = create_byte_stream(data) + result = csv_bytes_to_list(csv_bytes) + assert result == data + + +def test_create_byte_stream_empty_data(): + csv_bytes = create_byte_stream([]) + assert csv_bytes == b"" + + +def test_create_byte_stream_handles_quoted_fields(): + data = [ + {"student_id": "1234", "name": 'Student "One"', "course": "Course, A"}, + {"student_id": "5678", "name": 'Student "Two"', "course": "Course, B"}, + ] + csv_bytes = create_byte_stream(data) + result = csv_bytes_to_list(csv_bytes) + assert result == data + + +def test_create_byte_stream_consistent_header_order(): + data = [ + {"student_id": "1234", "name": "Alice", "course": "Math"}, + {"student_id": "5678", "name": "Bob", "course": "Science"}, + ] + csv_bytes = create_byte_stream(data) + csv_string = csv_bytes.decode("utf-8") + header_line = csv_string.splitlines()[0] + expected_header = ",".join(data[0].keys()) + assert header_line == expected_header + + +def test_create_byte_stream_special_characters(): + data = [ + {"student_id": "1234", "name": "Student 1", "course": "Line1\nLine2"}, + {"student_id": "5678", "name": "Student 2", "course": "Value with, comma"}, + ] + csv_bytes = create_byte_stream(data) + result = csv_bytes_to_list(csv_bytes) + assert result == data -- cgit v1.2.3 From ef3f16de8d93821d54344d5cdd16d8deee0b016c Mon Sep 17 00:00:00 2001 From: Alex Schofield Date: Wed, 19 Feb 2025 15:40:39 +0000 Subject: wrap write functions in class and update references --- cli.py | 6 ++++-- obfuscator/write.py | 24 ++++++++++++++---------- test/test_write.py | 2 +- 3 files changed, 19 insertions(+), 13 deletions(-) (limited to 'test/test_write.py') diff --git a/cli.py b/cli.py index 5100e2b..bd49707 100644 --- a/cli.py +++ b/cli.py @@ -1,8 +1,8 @@ import argparse from obfuscator.read import DataReader +from obfuscator.write import DataWriter from obfuscator.obfuscate import obfuscate from obfuscator.logger import get_logger -from obfuscator.csv_writer import create_byte_stream def main(): @@ -43,7 +43,9 @@ def main(): obfuscated_data = obfuscate(data, args.pii) - return create_byte_stream(obfuscated_data) + writer = DataWriter() + + return writer.create_byte_stream(obfuscated_data) if __name__ == "__main__": diff --git a/obfuscator/write.py b/obfuscator/write.py index de7cd4b..4081f0f 100644 --- a/obfuscator/write.py +++ b/obfuscator/write.py @@ -6,18 +6,22 @@ from obfuscator.logger import get_logger logger = get_logger("CSVWRITER") -def create_byte_stream(data: List[Dict[str, str]]) -> bytes: - if not data: - logger.error("Invalid or empty data was provided to write") +class DataWriter: + def __init__(self): + pass - output = io.StringIO() + def create_byte_stream(self, data: List[Dict[str, str]]) -> bytes: + if not data: + logger.error("Invalid or empty data was provided to write") - headers = list(data[0].keys()) + output = io.StringIO() - writer = csv.DictWriter(output, fieldnames=headers) - writer.writeheader() - writer.writerows(data) + headers = list(data[0].keys()) - csv_string = output.getvalue() + writer = csv.DictWriter(output, fieldnames=headers) + writer.writeheader() + writer.writerows(data) - return csv_string.encode("utf-8") + csv_string = output.getvalue() + + return csv_string.encode("utf-8") diff --git a/test/test_write.py b/test/test_write.py index eceac28..f339799 100644 --- a/test/test_write.py +++ b/test/test_write.py @@ -1,6 +1,6 @@ import io import csv -from obfuscator.csv_writer import create_byte_stream +from obfuscator.write import create_byte_stream def csv_bytes_to_list(csv_bytes: bytes): -- cgit v1.2.3 From ef204e7edabb7a3de2a747d83ae6a472b692fb38 Mon Sep 17 00:00:00 2001 From: Alex Schofield Date: Wed, 19 Feb 2025 15:43:08 +0000 Subject: partially fix tests by updating references with new class --- test/test_write.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) (limited to 'test/test_write.py') diff --git a/test/test_write.py b/test/test_write.py index f339799..4929b06 100644 --- a/test/test_write.py +++ b/test/test_write.py @@ -1,6 +1,8 @@ import io import csv -from obfuscator.write import create_byte_stream +from obfuscator.write import DataWriter + +writer = DataWriter() def csv_bytes_to_list(csv_bytes: bytes): @@ -15,13 +17,14 @@ def test_create_byte_stream_valid_data(): {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, ] - csv_bytes = create_byte_stream(data) + csv_bytes = writer.create_byte_stream(data) result = csv_bytes_to_list(csv_bytes) assert result == data def test_create_byte_stream_empty_data(): - csv_bytes = create_byte_stream([]) + data = [] + csv_bytes = writer.create_byte_stream(data) assert csv_bytes == b"" @@ -30,7 +33,7 @@ def test_create_byte_stream_handles_quoted_fields(): {"student_id": "1234", "name": 'Student "One"', "course": "Course, A"}, {"student_id": "5678", "name": 'Student "Two"', "course": "Course, B"}, ] - csv_bytes = create_byte_stream(data) + csv_bytes = writer.create_byte_stream(data) result = csv_bytes_to_list(csv_bytes) assert result == data @@ -40,7 +43,7 @@ def test_create_byte_stream_consistent_header_order(): {"student_id": "1234", "name": "Alice", "course": "Math"}, {"student_id": "5678", "name": "Bob", "course": "Science"}, ] - csv_bytes = create_byte_stream(data) + csv_bytes = writer.create_byte_stream(data) csv_string = csv_bytes.decode("utf-8") header_line = csv_string.splitlines()[0] expected_header = ",".join(data[0].keys()) @@ -52,6 +55,6 @@ def test_create_byte_stream_special_characters(): {"student_id": "1234", "name": "Student 1", "course": "Line1\nLine2"}, {"student_id": "5678", "name": "Student 2", "course": "Value with, comma"}, ] - csv_bytes = create_byte_stream(data) + csv_bytes = writer.create_byte_stream(data) result = csv_bytes_to_list(csv_bytes) assert result == data -- cgit v1.2.3