diff --git a/kartograf/context.py b/kartograf/context.py index 72e135c..1106c6c 100644 --- a/kartograf/context.py +++ b/kartograf/context.py @@ -1,5 +1,5 @@ from datetime import datetime -import os +from pathlib import Path import sys import time @@ -20,12 +20,12 @@ def __init__(self, args): self.reproduce = self.args.reproduce is not None if self.reproduce: - items = os.listdir(self.args.reproduce) + items = Path.iterdir(Path(self.args.reproduce)) source_folders = [] for folder in items: - full_path = os.path.join(self.args.reproduce, folder) - if os.path.isdir(full_path): - source_folders.append(folder) + full_path = self.args.reproduce / folder + if Path.is_dir(full_path): + source_folders.append(folder.name) # We override the args because we are reproducing and only where we # have data we try to use is, the actual args passed don't matter. self.args.irr = 'irr' in source_folders @@ -44,47 +44,52 @@ def __init__(self, args): self.epoch_dir = str(int(utc_time_now)) self.epoch_datetime = datetime.utcfromtimestamp(int(utc_time_now)) - cwd = os.getcwd() + cwd = Path.cwd() # Data dir if self.reproduce: if not self.args.reproduce.endswith('/'): self.args.reproduce += '/' self.data_dir = self.args.reproduce else: - self.data_dir = f"{cwd}/data/{self.epoch_dir}/" - self.data_dir_irr = f"{self.data_dir}irr/" - self.data_dir_rpki_cache = f"{self.data_dir}rpki/cache/" - self.data_dir_rpki_tals = f"{self.data_dir}rpki/tals/" - self.data_dir_collectors = f"{self.data_dir}collectors/" + self.data_dir = str(cwd / "data" / self.epoch_dir) + + if Path(self.data_dir).exists() and not self.reproduce: + print("Not so fast, a folder with that epoch already exists.") + sys.exit() + + self.data_dir_irr = str(Path(self.data_dir) / "irr") + self.data_dir_rpki_cache = str(Path(self.data_dir) / "rpki" / "cache") + self.data_dir_rpki_tals = str(Path(self.data_dir) / "rpki" / "tals") + self.data_dir_collectors = str(Path(self.data_dir) / "collectors") # Out dir - self.out_dir = f"{cwd}/out/{self.epoch_dir}/" - self.out_dir_irr = f"{self.out_dir}irr/" - self.out_dir_rpki = f"{self.out_dir}rpki/" - self.out_dir_collectors = f"{self.out_dir}collectors/" + self.out_dir = str(cwd / "out" / self.epoch_dir) + self.out_dir_irr = str(Path(self.out_dir) / "irr") + self.out_dir_rpki = str(Path(self.out_dir) / "rpki") + self.out_dir_collectors = str(Path(self.out_dir) / "collectors") - if os.path.exists(self.data_dir) and not self.reproduce: + if Path(self.data_dir).exists() and not self.reproduce: print("Not so fast, a folder with that epoch already exists.") sys.exit() # We skip creating the folders if we are reproducing a run. if not self.reproduce: - os.makedirs(self.data_dir_rpki_cache) - os.makedirs(self.data_dir_rpki_tals) + Path(self.data_dir_rpki_cache).mkdir(parents=True) + Path(self.data_dir_rpki_tals).mkdir(parents=True) if self.args.irr: - os.makedirs(self.data_dir_irr) + Path(self.data_dir_irr).mkdir(parents=True) if self.args.routeviews: - os.makedirs(self.data_dir_collectors) - os.makedirs(self.out_dir_rpki) + Path(self.data_dir_collectors).mkdir(parents=True) + Path(self.out_dir_rpki).mkdir(parents=True) if self.args.irr: - os.makedirs(self.out_dir_irr) + Path(self.out_dir_irr).mkdir(parents=True) if self.args.routeviews: - os.makedirs(self.out_dir_collectors) + Path(self.out_dir_collectors).mkdir(parents=True) - self.final_result_file = f"{self.out_dir}final_result.txt" + self.final_result_file = str(Path(self.out_dir) / "final_result.txt") self.max_encode = self.args.max_encode if self.args.debug: - self.debug_log = f"{self.out_dir}debug.log" + self.debug_log = str(Path(self.out_dir) / "debug.log") else: self.debug_log = "" diff --git a/kartograf/rpki/parse.py b/kartograf/rpki/parse.py index 49c035a..2fc0a72 100644 --- a/kartograf/rpki/parse.py +++ b/kartograf/rpki/parse.py @@ -1,4 +1,5 @@ import json +from pathlib import Path from typing import Dict from kartograf.bogon import ( @@ -12,8 +13,8 @@ @timed def parse_rpki(context): - raw_input = f"{context.out_dir_rpki}rpki_raw.json" - rpki_res = f"{context.out_dir_rpki}rpki_final.txt" + raw_input = Path(context.out_dir_rpki) / "rpki_raw.json" + rpki_res = Path(context.out_dir_rpki) / "rpki_final.txt" output_cache: Dict[str, [str, str]] = {} diff --git a/tests/test_context.py b/tests/test_context.py index d7dadd8..f9196e3 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -1,4 +1,5 @@ import os +from pathlib import Path import pytest from kartograf.context import Context from kartograf.cli import create_parser @@ -21,7 +22,7 @@ def test_basic_map_context(parser, tmp_path): assert isinstance(context.epoch, str) assert isinstance(int(context.epoch), int) assert context.max_encode == 33521664 - assert context.debug_log.endswith('debug.log') + assert Path(context.debug_log).name == 'debug.log' def test_map_context_with_reproduce(parser, tmp_path): # Setup a mock reproduction directory @@ -53,10 +54,17 @@ def test_directory_creation(parser, tmp_path): os.chdir(tmp_path) context = Context(args) - assert os.path.exists(context.data_dir_rpki_cache) - assert os.path.exists(context.data_dir_rpki_tals) - assert os.path.exists(context.data_dir_irr) - assert os.path.exists(context.data_dir_collectors) - assert os.path.exists(context.out_dir_rpki) - assert os.path.exists(context.out_dir_irr) - assert os.path.exists(context.out_dir_collectors) + assert isinstance(context.data_dir_rpki_cache, str) + assert Path(context.data_dir_rpki_cache).exists() + assert isinstance(context.data_dir_rpki_tals, str) + assert Path(context.data_dir_rpki_tals).exists() + assert isinstance(context.data_dir_irr, str) + assert Path(context.data_dir_irr).exists() + assert isinstance(context.data_dir_collectors, str) + assert Path(context.data_dir_collectors).exists() + assert isinstance(context.out_dir_rpki, str) + assert Path(context.out_dir_rpki).exists() + assert isinstance(context.out_dir_irr, str) + assert Path(context.out_dir_irr).exists() + assert isinstance(context.out_dir_collectors, str) + assert Path(context.out_dir_collectors).exists() diff --git a/tests/test_merge.py b/tests/test_merge.py index 70a9528..bc455fc 100644 --- a/tests/test_merge.py +++ b/tests/test_merge.py @@ -40,10 +40,10 @@ def test_merge_from_fixtures(tmp_path): and validates against expected networks, i.e. subnets are merged into the root network appropriately. ''' - testdir = Path(__file__).parent - base_nets, base_results = __read_test_vectors(testdir / "data/base_file.csv") + data_dir = Path(__file__).parent / "data" + base_nets, base_results = __read_test_vectors(data_dir / "base_file.csv") base_path = tmp_path / "base.txt" - extra_nets, extra_results = __read_test_vectors(testdir / "data/extra_file.csv") + extra_nets, extra_results = __read_test_vectors(data_dir / "extra_file.csv") extra_path = tmp_path / "extra.txt" # write the networks to disk, generating ASNs for each network generate_ip_file(base_path, build_file_lines(base_nets, generate_asns(len(base_nets))))