diff --git a/src/audit_manager.py b/src/audit_manager.py index b90dc993..2210416b 100644 --- a/src/audit_manager.py +++ b/src/audit_manager.py @@ -137,18 +137,17 @@ def test_for_anti_bot(self) -> str: # Write to anti-bot.csv csv_writer = CSVWriter() - csv_writer.add_rows( - [ - { - 'organisation': org['organisation'], - 'domain': netloc, - 'url': url, - 'anti_bot_check': status, - 'viewport_size': self.browser.viewport_size, - } - ] + csv_writer.append_rows( + f'./results/{self.config.audit_name}/anti_bot.csv', + { + 'organisation': org['organisation'], + 'domain': netloc, + 'url': url, + 'anti_bot_check': status, + 'viewport_size': self.browser.viewport_size, + }, ) - csv_writer.write_csv_file(f'./results/{self.config.audit_name}/anti_bot.csv') + self.discarded_urls[url] = status return status @@ -287,7 +286,7 @@ def run_audits(self) -> bool: # noqa: PLR0915 test_instance = audit['audit_class'](config=self.config, browser=self.browser, **audit['kwargs']) try: - audit_result = test_instance.run() + audit_result: list[dict[str, Any]] | bool = test_instance.run() except selenium.common.exceptions.WebDriverException: logger.exception( 'Due to WebDriverException, test %s skipped on viewport %s for website %s', @@ -359,8 +358,7 @@ def run_audits(self) -> bool: # noqa: PLR0915 # Write results csv_writer = CSVWriter() - csv_writer.add_rows(audit_result) - csv_writer.write_csv_file(f'./results/{self.config.audit_name}/{audit_name}.csv') + csv_writer.append_rows(f'./results/{self.config.audit_name}/{audit_name}.csv', *audit_result) # At least one audit successfully produced results any_audit_succeeded = True diff --git a/src/crawler.py b/src/crawler.py index fff7389e..6a00c286 100644 --- a/src/crawler.py +++ b/src/crawler.py @@ -366,18 +366,17 @@ def are_url_headers_acceptable(self, base_url: str, parent_url: str, url_data: s url_data['final_url'], url_data['status_code'], ) - # Write bad response codes with CSVWriter - csv_writer = src.output.CSVWriter() - csv_writer.add_row( - { - 'base_url': base_url, - 'parent_url': parent_url, - 'url': url_data['final_url'], - 'status_code': url_data['status_code'], - } - ) if self.config.record_unexpected_response_codes: - csv_writer.write_csv_file(f'./results/{self.config.audit_name}/unexpected_response_codes.csv') + csv_writer = src.output.CSVWriter() + csv_writer.append_rows( + f'./results/{self.config.audit_name}/unexpected_response_codes.csv', + { + 'base_url': base_url, + 'parent_url': parent_url, + 'url': url_data['final_url'], + 'status_code': url_data['status_code'], + }, + ) return False return src.filters.url_filter_by_header_content_type(url_data['final_url'], url_data['headers']) @@ -572,17 +571,15 @@ def crawl(self, site_data: SiteData, base_url: str) -> None: # noqa: PLR0912, P # Write to audit_log.csv csv_writer = CSVWriter() - csv_writer.add_rows( - [ - { - 'organisation': site_data['organisation'], - 'base_url': site_data['url'], - 'url': url, - 'sector': site_data['sector'], - } - ] + csv_writer.append_rows( + f'./results/{self.config.audit_name}/audit_log.csv', + { + 'organisation': site_data['organisation'], + 'base_url': site_data['url'], + 'url': url, + 'sector': site_data['sector'], + }, ) - csv_writer.write_csv_file(f'./results/{self.config.audit_name}/audit_log.csv') self.register_audit_plugins(audit_manager, url, site_data) test_success = audit_manager.run_audits() @@ -649,17 +646,15 @@ def record_pages_scanned(self, site_data: SiteData, pages_scanned: int) -> None: """Record the number of pages that were scanned for the site.""" with self.config.lock: csv_writer = src.output.CSVWriter() - csv_writer.add_rows( - [ - { - 'organisation': site_data['organisation'], - 'base_url': site_data['url'], - 'number_of_pages': pages_scanned, - 'sector': site_data['sector'], - } - ] + csv_writer.append_rows( + f'./results/{self.config.audit_name}/pages_scanned.csv', + { + 'organisation': site_data['organisation'], + 'base_url': site_data['url'], + 'number_of_pages': pages_scanned, + 'sector': site_data['sector'], + }, ) - csv_writer.write_csv_file(f'./results/{self.config.audit_name}/pages_scanned.csv') class RandomQueue[T]: diff --git a/src/output.py b/src/output.py index aaae7ddb..cf2ffb07 100644 --- a/src/output.py +++ b/src/output.py @@ -11,8 +11,6 @@ from config import Config -# pylint: disable=too-many-locals - logger = logging.getLogger('cwac') @@ -26,11 +24,7 @@ class CSVWriter: # A lock to prevent multiple threads writing to file_locks dict lock_for_file_locks = threading.Lock() - def __init__(self) -> None: - """Init variables.""" - self.rows: list[dict[Any, Any]] = [] - - def get_file_lock(self, path: str) -> threading.Lock: + def _get_file_lock(self, path: str) -> threading.Lock: """Get a lock for a file. Args: @@ -44,64 +38,25 @@ def get_file_lock(self, path: str) -> threading.Lock: CSVWriter.file_locks[path] = threading.Lock() return CSVWriter.file_locks[path] - def read_csv(self, path: str) -> list[dict[Any, Any]]: - """Read a CSV file as a list of dictionaries. + def append_rows(self, path: str, *rows: dict[Any, Any]) -> None: + """Append one or more rows to a CSV file. Args: - path (str): path to CSV file - - Returns: - list[dict[Any, Any]]: list of dictionaries - """ - with self.get_file_lock(path), open(path, encoding='utf-8-sig') as csvfile: - reader = csv.DictReader(csvfile) - rows = list(reader) - return rows - - def add_row(self, row: dict[Any, Any]) -> None: - """Add a row to the CSV row buffer. - - Args: - row (dict[Any, Any]): A dictionary of row contents - """ - self.rows.append(row) - - def add_rows(self, rows: list[dict[Any, Any]]) -> None: - """Add a list of rows to the CSV row buffer. - - Args: - rows (list[dict[Any, Any]]): list of rows of data - """ - for row in rows: - self.rows.append(row) - - def write_csv_file(self, path: str, overwrite: bool = False) -> bool: - """Write data to a CSV file. - - Args: - path (str): path to write data - overwrite (bool): overwrite existing file - - Returns: - bool: True if write successful, else False + path (str): path to file + *rows (dict[str, Any]): one or more rows of data """ - if not self.rows: - return False + if not rows: + return - keys = self.rows[0].keys() + keys = rows[0].keys() - with self.get_file_lock(path): - file_exists = False if overwrite else os.path.exists(path) - file_mode = 'w' if overwrite else 'a' - with open(path, file_mode, encoding='utf-8-sig') as csvfile: + with self._get_file_lock(path): + file_already_exists = os.path.exists(path) + with open(path, 'a', encoding='utf-8-sig') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=keys) - if not file_exists: + if not file_already_exists: writer.writeheader() - writer.writerows(self.rows) - - self.rows = [] - - return True + writer.writerows(rows) def output_init_message(config: Config) -> None: @@ -212,20 +167,18 @@ def print_progress_bar( # Write progress data to CSV file csv_writer = CSVWriter() - - output_row = { - 'time': time.time(), - 'iteration': iteration, - 'total': total, - 'speed': f'{speed:.2f}', - 'percent': percent, - 'elapsed': f'{elapsed}', - 'remaining': f'{time_est}', - } - - csv_writer.add_row(output_row) - - csv_writer.write_csv_file(f'./results/{config.audit_name}/progress.csv') + csv_writer.append_rows( + f'./results/{config.audit_name}/progress.csv', + { + 'time': time.time(), + 'iteration': iteration, + 'total': total, + 'speed': f'{speed:.2f}', + 'percent': percent, + 'elapsed': f'{elapsed}', + 'remaining': f'{time_est}', + }, + ) # Print New Line on Complete if iteration == total: diff --git a/tests/test_output.py b/tests/test_output.py new file mode 100644 index 00000000..11adabb8 --- /dev/null +++ b/tests/test_output.py @@ -0,0 +1,155 @@ +"""Tests the behaviour of the output classes and functions.""" + +import os +import textwrap + +import pytest + +from src.output import CSVWriter + + +@pytest.mark.usefixtures('fs') +class TestCSVWriter: + """Tests writing csv files using the CSVWriter class.""" + + def test_csv_is_not_written_when_no_rows(self) -> None: + """Skips writing a file when no rows have been added.""" + writer = CSVWriter() + + writer.append_rows('file.csv') + + assert os.path.exists('file.csv') is not True + + def test_csv_is_written(self) -> None: + """Writes each row to a csv file, with a header.""" + writer = CSVWriter() + + writer.append_rows( + 'file.csv', + {'name': 'Bob', 'age': 20, 'location': 'Wellington'}, + {'name': 'Alice', 'age': 31, 'location': 'Wellington'}, + {'name': 'Greg', 'age': 23, 'location': 'Auckland'}, + ) + + assert os.path.exists('file.csv') is True + + expected = """ + name,age,location + Bob,20,Wellington + Alice,31,Wellington + Greg,23,Auckland + """ + + with open('file.csv', encoding='utf-8-sig') as f: + assert f.read() == textwrap.dedent(expected).lstrip() + + def test_csv_is_written_with_bom(self) -> None: + """Writes a bom before the header row.""" + writer = CSVWriter() + + writer.append_rows( + 'file.csv', + {'name': 'Bob', 'age': 20, 'location': 'Wellington'}, + {'name': 'Alice', 'age': 31, 'location': 'Wellington'}, + {'name': 'Greg', 'age': 23, 'location': 'Auckland'}, + ) + + assert os.path.exists('file.csv') is True + + expected = """ + \ufeffname,age,location + Bob,20,Wellington + Alice,31,Wellington + Greg,23,Auckland + """ + + with open('file.csv', encoding='utf-8') as f: + assert f.read() == textwrap.dedent(expected).lstrip() + + def test_appends_to_existing_file(self) -> None: + """Writes new rows to a csv file, without duplicating the header.""" + writer = CSVWriter() + + writer.append_rows( + 'file.csv', + {'name': 'Bob', 'age': 20, 'location': 'Wellington'}, + ) + + assert os.path.exists('file.csv') is True + + writer.append_rows( + 'file.csv', + {'name': 'Alice', 'age': 31, 'location': 'Wellington'}, + {'name': 'Greg', 'age': 23, 'location': 'Auckland'}, + ) + + expected = """ + name,age,location + Bob,20,Wellington + Alice,31,Wellington + Greg,23,Auckland + """ + + with open('file.csv', encoding='utf-8-sig') as f: + assert f.read() == textwrap.dedent(expected).lstrip() + + def test_ignores_missing_columns(self) -> None: + """Ignores columns missing in subsequent rows.""" + writer = CSVWriter() + + writer.append_rows( + 'file.csv', + {'name': 'Bob', 'age': 20, 'location': 'Wellington'}, + {'name': 'Alice', 'age': 31}, + {'age': 23, 'location': 'Auckland'}, + ) + + assert os.path.exists('file.csv') is True + + expected = """ + name,age,location + Bob,20,Wellington + Alice,31, + ,23,Auckland + """ + + with open('file.csv', encoding='utf-8-sig') as f: + assert f.read() == textwrap.dedent(expected).lstrip() + + def test_raises_extra_columns(self) -> None: + """Errors when rows have extra columns.""" + writer = CSVWriter() + + with pytest.raises(ValueError): + writer.append_rows( + 'file.csv', + {'name': 'Bob', 'age': 20}, + {'name': 'Alice', 'age': 31, 'location': 'Wellington'}, + {'name': 'Greg', 'age': 23, 'location': 'Auckland'}, + ) + + # the file will still end up being created due to the open mode + assert os.path.exists('file.csv') + + def test_column_ordering_is_consistent(self) -> None: + """Orders columns based on the first row.""" + writer = CSVWriter() + + writer.append_rows( + 'file.csv', + {'name': 'Bob', 'age': 20, 'location': 'Wellington'}, + {'age': 31, 'name': 'Alice', 'location': 'Wellington'}, + {'location': 'Auckland', 'age': 23, 'name': 'Greg'}, + ) + + assert os.path.exists('file.csv') is True + + expected = """ + name,age,location + Bob,20,Wellington + Alice,31,Wellington + Greg,23,Auckland + """ + + with open('file.csv', encoding='utf-8-sig') as f: + assert f.read() == textwrap.dedent(expected).lstrip()