From ae89b05dbc8feebc1410f39143c0d829f8704235 Mon Sep 17 00:00:00 2001 From: Alex Schofield Date: Wed, 19 Feb 2025 15:37:16 +0000 Subject: rename CSVReader to DataReader & update references --- cli.py | 4 +- obfuscator/csv_reader.py | 89 ----------------------- obfuscator/read.py | 89 +++++++++++++++++++++++ test/test_csv_reader.py | 181 ----------------------------------------------- test/test_read.py | 181 +++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 272 insertions(+), 272 deletions(-) delete mode 100644 obfuscator/csv_reader.py create mode 100644 obfuscator/read.py delete mode 100644 test/test_csv_reader.py create mode 100644 test/test_read.py diff --git a/cli.py b/cli.py index 7ffccd8..5100e2b 100644 --- a/cli.py +++ b/cli.py @@ -1,5 +1,5 @@ import argparse -from obfuscator.csv_reader import CSVReader +from obfuscator.read import DataReader from obfuscator.obfuscate import obfuscate from obfuscator.logger import get_logger from obfuscator.csv_writer import create_byte_stream @@ -32,7 +32,7 @@ def main(): log_level = "DEBUG" if args.verbose else "INFO" logger = get_logger("CLI", log_level) - reader = CSVReader(log_level) + reader = DataReader(log_level) if args.local and not args.s3: logger.debug("Read data from local path") diff --git a/obfuscator/csv_reader.py b/obfuscator/csv_reader.py deleted file mode 100644 index 1f503d7..0000000 --- a/obfuscator/csv_reader.py +++ /dev/null @@ -1,89 +0,0 @@ -import csv -import io -import boto3 -import os -from typing import List, Dict -from obfuscator.logger import get_logger -from obfuscator.utils import Utilities - - -class CSVReader: - """ - A class to read CSV data from a local file, S3 object, or string. Near - the project completion, support for JSON/Parquet files will be added. - """ - - def __init__(self, log_level=None): - self.log_level = log_level - self.logger = get_logger("CSVREADER", log_level) - - def read_local(self, path) -> List[Dict[str, str]]: - """ - A method to read a local CSV file and return the data as a list of - dictionaries. - """ - self.logger.debug(f"Reading local CSV from: {path}") - - try: - with open(path, mode="r", encoding="utf-8") as f: - reader = csv.DictReader(f) - return [dict(row) for row in reader] - except FileNotFoundError: - self.logger.error(f"File not found: {path}") - raise - except Exception as e: - self.logger.error(f"Error reading file: {e}") - - def read_s3(self, path) -> List[Dict[str, str]]: - """ - A method to read an S3 object containing CSV data - and return the data as a list of dictionaries. - """ - utils = Utilities(self.log_level) - bucket, key = utils.get_s3_path(path) - self.logger.debug(f"Reading S3 CSV from: {bucket}/{key}") - - if os.getenv("LOCALSTACK", "FALSE").upper() == "TRUE": - localstack_endpoint = "http://localhost.localstack.cloud:4566" - self.logger.debug( - "Using LocalStack endpoint for S3 - ensure LocalStack is running" - ) - client = boto3.client( - "s3", - endpoint_url=localstack_endpoint, - aws_access_key_id="dummy", - aws_secret_access_key="dummy", - ) - self.logger.debug(f"endpoint_url: {localstack_endpoint}") - else: - client = boto3.client("s3") - - try: - response = client.get_object(Bucket=bucket, Key=key) - self.logger.info("S3 object read successfully") - content = response["Body"].read().decode("utf-8") - return self.read_string(content) - except client.exceptions.NoSuchKey: - self.logger.error(f"Object not found: {bucket}/{key}") - raise - except client.exceptions.ClientError as e: - self.logger.error(f"Error reading S3 object: {e}") - raise - except UnicodeDecodeError as e: - self.logger.error(f"Error decoding S3 object: {e}") - raise - except Exception as e: - self.logger.error(f"Error reading S3 object: {e}") - raise - - def read_string(self, content: str) -> List[Dict[str, str]]: - """ - A method to read CSV data from a string and return the data as a list - of dictionaries. - """ - if not content.strip(): - return [] - - f = io.StringIO(content) - reader = csv.DictReader(f) - return [dict(row) for row in reader] diff --git a/obfuscator/read.py b/obfuscator/read.py new file mode 100644 index 0000000..b704643 --- /dev/null +++ b/obfuscator/read.py @@ -0,0 +1,89 @@ +import csv +import io +import boto3 +import os +from typing import List, Dict +from obfuscator.logger import get_logger +from obfuscator.utils import Utilities + + +class DataReader: + """ + A class to read CSV data from a local file, S3 object, or string. Near + the project completion, support for JSON/Parquet files will be added. + """ + + def __init__(self, log_level=None): + self.log_level = log_level + self.logger = get_logger("CSVREADER", log_level) + + def read_local(self, path) -> List[Dict[str, str]]: + """ + A method to read a local CSV file and return the data as a list of + dictionaries. + """ + self.logger.debug(f"Reading local CSV from: {path}") + + try: + with open(path, mode="r", encoding="utf-8") as f: + reader = csv.DictReader(f) + return [dict(row) for row in reader] + except FileNotFoundError: + self.logger.error(f"File not found: {path}") + raise + except Exception as e: + self.logger.error(f"Error reading file: {e}") + + def read_s3(self, path) -> List[Dict[str, str]]: + """ + A method to read an S3 object containing CSV data + and return the data as a list of dictionaries. + """ + utils = Utilities(self.log_level) + bucket, key = utils.get_s3_path(path) + self.logger.debug(f"Reading S3 CSV from: {bucket}/{key}") + + if os.getenv("LOCALSTACK", "FALSE").upper() == "TRUE": + localstack_endpoint = "http://localhost.localstack.cloud:4566" + self.logger.debug( + "Using LocalStack endpoint for S3 - ensure LocalStack is running" + ) + client = boto3.client( + "s3", + endpoint_url=localstack_endpoint, + aws_access_key_id="dummy", + aws_secret_access_key="dummy", + ) + self.logger.debug(f"endpoint_url: {localstack_endpoint}") + else: + client = boto3.client("s3") + + try: + response = client.get_object(Bucket=bucket, Key=key) + self.logger.info("S3 object read successfully") + content = response["Body"].read().decode("utf-8") + return self.read_string(content) + except client.exceptions.NoSuchKey: + self.logger.error(f"Object not found: {bucket}/{key}") + raise + except client.exceptions.ClientError as e: + self.logger.error(f"Error reading S3 object: {e}") + raise + except UnicodeDecodeError as e: + self.logger.error(f"Error decoding S3 object: {e}") + raise + except Exception as e: + self.logger.error(f"Error reading S3 object: {e}") + raise + + def read_string(self, content: str) -> List[Dict[str, str]]: + """ + A method to read CSV data from a string and return the data as a list + of dictionaries. + """ + if not content.strip(): + return [] + + f = io.StringIO(content) + reader = csv.DictReader(f) + return [dict(row) for row in reader] diff --git a/test/test_csv_reader.py b/test/test_csv_reader.py deleted file mode 100644 index d6129e7..0000000 --- a/test/test_csv_reader.py +++ /dev/null @@ -1,181 +0,0 @@ -import boto3 -from moto import mock_aws -from obfuscator.csv_reader import CSVReader -import pytest - -reader = CSVReader(log_level="DEBUG") - - -def test_empty_csv_should_return_no_content(): - content = "" - result = reader.read_string(content) - expected = [] - assert result == expected - - -def test_csv_with_header_only_should_return_no_content(): - content = "student_id,name,course\n" - result = reader.read_string(content) - expected = [] - assert result == expected - - -def test_csv_with_valid_data(): - content = ( - "student_id,name,course\n" - "1234,Student 1,Course 1\n" - "5678,Student 2,Course 2\n" - ) - result = reader.read_string(content) - expected = [ - {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, - {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, - ] - assert result == expected - - -def test_csv_with_quoted_fields_should_run_as_expected(): - content = ( - "student_id,name,course\n" - '1234,"Student 1","Course 1"\n' - '5678,"Student 2","Course 2"\n' - ) - result = reader.read_string(content) - expected = [ - {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, - {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, - ] - assert result == expected - - -def setup_s3(s3_client, bucket: str, key: str, content: str): - s3_client.create_bucket( - Bucket=bucket, - CreateBucketConfiguration={"LocationConstraint": "eu-west-2"}, - ) - s3_client.put_object(Bucket=bucket, Key=key, Body=content) - - -@pytest.fixture(autouse=True) -def s3_client(): - with mock_aws(): - yield boto3.client("s3", "eu-west-2") - - -def test_read_s3_valid_csv_returns_expected(): - with mock_aws(): - s3 = boto3.client("s3", region_name="eu-west-2") - bucket = "test-bucket" - key = "data/mock.csv" - - csv_content = ( - "student_id,name,course\n" - "1234,Student 1,Course 1\n" - "5678,Student 2,Course 2\n" - ) - - setup_s3(s3, bucket, key, csv_content) - path = f"s3://{bucket}/{key}" - - data = reader.read_s3(path) - - expected = [ - {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, - {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, - ] - - assert data == expected - - -def test_read_s3_empty_csv_returns_empty_list(): - with mock_aws(): - s3 = boto3.client("s3", region_name="eu-west-2") - bucket = "empty-bucket" - key = "data/empty.csv" - csv_content = "student_id,name,course\n" - setup_s3(s3, bucket, key, csv_content) - path = f"s3://{bucket}/{key}" - - data = reader.read_s3(path) - assert data == [] - - -def test_read_s3_nonexistent_bucket_raises_exception(): - with mock_aws(): - bucket = "nonexistent-bucket" - key = "data/mock.csv" - path = f"s3://{bucket}/{key}" - with pytest.raises(Exception): - reader.read_s3(path) - - -def test_read_s3_nonexistent_key_raises_exception(): - with mock_aws(): - s3 = boto3.client("s3", region_name="eu-west-2") - bucket = "test-bucket" - s3.create_bucket( - Bucket=bucket, - CreateBucketConfiguration={"LocationConstraint": "eu-west-2"}, - ) - key = "data/nonexistent.csv" - path = f"s3://{bucket}/{key}" - with pytest.raises(Exception): - reader.read_s3(path) - - -def test_read_s3_malformed_csv_returns_expected(): - with mock_aws(): - s3 = boto3.client("s3", region_name="eu-west-2") - bucket = "test-bucket" - key = "data/malformed.csv" - csv_content = "1234,Student 1,Course 1\n" "5678,Student 2,Course 2\n" - setup_s3(s3, bucket, key, csv_content) - path = f"s3://{bucket}/{key}" - - data = reader.read_s3(path) - expected = [{"1234": "5678", "Student 1": "Student 2", "Course 1": "Course 2"}] - assert data == expected - - -def test_read_s3_csv_with_extra_empty_lines(): - with mock_aws(): - s3 = boto3.client("s3", region_name="eu-west-2") - bucket = "test-bucket" - key = "data/extra_lines.csv" - csv_content = ( - "student_id,name,course\n" - "1234,Student 1,Course 1\n" - "\n" - "5678,Student 2,Course 2\n" - "\n" - ) - setup_s3(s3, bucket, key, csv_content) - path = f"s3://{bucket}/{key}" - - data = reader.read_s3(path) - expected = [ - {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, - {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, - ] - assert data == expected - - -def test_read_s3_csv_with_whitespace_in_fields(): - with mock_aws(): - s3 = boto3.client("s3", region_name="eu-west-2") - bucket = "test-bucket" - key = "data/whitespace.csv" - csv_content = ( - "student_id, name , course \n" - " 1234 , Student 1 , Course 1 \n" - "5678,Student 2,Course 2\n" - ) - setup_s3(s3, bucket, key, csv_content) - path = f"s3://{bucket}/{key}" - - data = reader.read_s3(path) - expected = [ - {"student_id": " 1234 ", " name ": " Student 1 ", " course ": " Course 1 "}, - {"student_id": "5678", " name ": "Student 2", " course ": "Course 2"}, - ] - assert data == expected diff --git a/test/test_read.py b/test/test_read.py new file mode 100644 index 0000000..903ab5d --- /dev/null +++ b/test/test_read.py @@ -0,0 +1,181 @@ +import boto3 +from moto import mock_aws +from obfuscator.read import CSVReader +import pytest + +reader = CSVReader(log_level="DEBUG") + + +def test_empty_csv_should_return_no_content(): + content = "" + result = reader.read_string(content) + expected = [] + assert result == expected + + +def test_csv_with_header_only_should_return_no_content(): + content = "student_id,name,course\n" + result = reader.read_string(content) + expected = [] + assert result == expected + + +def test_csv_with_valid_data(): + content = ( + "student_id,name,course\n" + "1234,Student 1,Course 1\n" + "5678,Student 2,Course 2\n" + ) + result = reader.read_string(content) + expected = [ + {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, + {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, + ] + assert result == expected + + +def test_csv_with_quoted_fields_should_run_as_expected(): + content = ( + "student_id,name,course\n" + '1234,"Student 1","Course 1"\n' + '5678,"Student 2","Course 2"\n' + ) + result = reader.read_string(content) + expected = [ + {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, + {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, + ] + assert result == expected + + +def setup_s3(s3_client, bucket: str, key: str, content: str): + s3_client.create_bucket( + Bucket=bucket, + CreateBucketConfiguration={"LocationConstraint": "eu-west-2"}, + ) + s3_client.put_object(Bucket=bucket, Key=key, Body=content) + + +@pytest.fixture(autouse=True) +def s3_client(): + with mock_aws(): + yield boto3.client("s3", "eu-west-2") + + +def test_read_s3_valid_csv_returns_expected(): + with mock_aws(): + s3 = boto3.client("s3", region_name="eu-west-2") + bucket = "test-bucket" + key = "data/mock.csv" + + csv_content = ( + "student_id,name,course\n" + "1234,Student 1,Course 1\n" + "5678,Student 2,Course 2\n" + ) + + setup_s3(s3, bucket, key, csv_content) + path = f"s3://{bucket}/{key}" + + data = reader.read_s3(path) + + expected = [ + {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, + {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, + ] + + assert data == expected + + +def test_read_s3_empty_csv_returns_empty_list(): + with mock_aws(): + s3 = boto3.client("s3", region_name="eu-west-2") + bucket = "empty-bucket" + key = "data/empty.csv" + csv_content = "student_id,name,course\n" + setup_s3(s3, bucket, key, csv_content) + path = f"s3://{bucket}/{key}" + + data = reader.read_s3(path) + assert data == [] + + +def test_read_s3_nonexistent_bucket_raises_exception(): + with mock_aws(): + bucket = "nonexistent-bucket" + key = "data/mock.csv" + path = f"s3://{bucket}/{key}" + with pytest.raises(Exception): + reader.read_s3(path) + + +def test_read_s3_nonexistent_key_raises_exception(): + with mock_aws(): + s3 = boto3.client("s3", region_name="eu-west-2") + bucket = "test-bucket" + s3.create_bucket( + Bucket=bucket, + CreateBucketConfiguration={"LocationConstraint": "eu-west-2"}, + ) + key = "data/nonexistent.csv" + path = f"s3://{bucket}/{key}" + with pytest.raises(Exception): + reader.read_s3(path) + + +def test_read_s3_malformed_csv_returns_expected(): + with mock_aws(): + s3 = boto3.client("s3", region_name="eu-west-2") + bucket = "test-bucket" + key = "data/malformed.csv" + csv_content = "1234,Student 1,Course 1\n" "5678,Student 2,Course 2\n" + setup_s3(s3, bucket, key, csv_content) + path = f"s3://{bucket}/{key}" + + data = reader.read_s3(path) + expected = [{"1234": "5678", "Student 1": "Student 2", "Course 1": "Course 2"}] + assert data == expected + + +def test_read_s3_csv_with_extra_empty_lines(): + with mock_aws(): + s3 = boto3.client("s3", region_name="eu-west-2") + bucket = "test-bucket" + key = "data/extra_lines.csv" + csv_content = ( + "student_id,name,course\n" + "1234,Student 1,Course 1\n" + "\n" + "5678,Student 2,Course 2\n" + "\n" + ) + setup_s3(s3, bucket, key, csv_content) + path = f"s3://{bucket}/{key}" + + data = reader.read_s3(path) + expected = [ + {"student_id": "1234", "name": "Student 1", "course": "Course 1"}, + {"student_id": "5678", "name": "Student 2", "course": "Course 2"}, + ] + assert data == expected + + +def test_read_s3_csv_with_whitespace_in_fields(): + with mock_aws(): + s3 = boto3.client("s3", region_name="eu-west-2") + bucket = "test-bucket" + key = "data/whitespace.csv" + csv_content = ( + "student_id, name , course \n" + " 1234 , Student 1 , Course 1 \n" + "5678,Student 2,Course 2\n" + ) + setup_s3(s3, bucket, key, csv_content) + path = f"s3://{bucket}/{key}" + + data = reader.read_s3(path) + expected = [ + {"student_id": " 1234 ", " name ": " Student 1 ", " course ": " Course 1 "}, + {"student_id": "5678", " name ": "Student 2", " course ": "Course 2"}, + ] + assert data == expected -- cgit v1.2.3