diff --git a/src/collect_data.py b/src/collect_data.py index 5988af6..867ab73 100644 --- a/src/collect_data.py +++ b/src/collect_data.py @@ -10,8 +10,8 @@ import google.cloud.bigquery as bq import json import argparse +import logging from yaml import safe_load -from time import time from helpers.helper import ROOT_DIR, INPUT_DIR @@ -29,32 +29,30 @@ def collect_data(force_query): filename = f'{ledger}_raw_data.json' file = INPUT_DIR / filename if not force_query and file.is_file(): - print(f'{ledger} data already exists locally. ' - f'For querying {ledger} anyway please run the script using the flag --force-query') + logging.info(f'{ledger} data already exists locally. ' + f'For querying {ledger} anyway please run the script using the flag --force-query') continue - print(f"Querying {ledger}..") - start = time() + logging.info(f"Querying {ledger}..") query = (queries[ledger]) query_job = client.query(query) try: rows = query_job.result() - print(f'Done querying {ledger} (took about {round(time() - start)} seconds)') + logging.info(f'Done querying {ledger}') except Exception as e: - print(f'{ledger} query failed, please make sure it is properly defined.') - print(f'The following exception was raised: {repr(e)}') + logging.info(f'{ledger} query failed, please make sure it is properly defined.') + logging.info(f'The following exception was raised: {repr(e)}') continue - print(f"Writing {ledger} data to file..") - start = time() + logging.info(f"Writing {ledger} data to file..") # write json lines to file with open(file, 'w') as f: for row in rows: f.write(json.dumps(dict(row), default=str) + "\n") - print(f'Done writing {ledger} data (took about {round(time() - start)} seconds)') - print(50 * '-') + logging.info(f'Done writing {ledger} data to file.\n') if __name__ == '__main__': + logging.basicConfig(format='[%(asctime)s] %(message)s', datefmt='%Y/%m/%d %I:%M:%S %p', level=logging.INFO) parser = argparse.ArgumentParser() parser.add_argument( '--force-query', diff --git a/tests/test_2_helper.py b/tests/test_2_helper.py index d15d177..ee3e746 100644 --- a/tests/test_2_helper.py +++ b/tests/test_2_helper.py @@ -16,12 +16,12 @@ def setup_and_cleanup(): The part before the yield command is run before the test (setup) and the part after the yield command is run after (cleanup) """ - print("Setting up") + # Setting up test_output_dir = OUTPUT_DIR / "test_output" if not os.path.exists(test_output_dir): os.makedirs(test_output_dir) yield test_output_dir - print("Cleaning up") + # Cleaning up shutil.rmtree(test_output_dir) diff --git a/tests/test_3_parsers.py b/tests/test_3_parsers.py index 33fbf66..423d5de 100644 --- a/tests/test_3_parsers.py +++ b/tests/test_3_parsers.py @@ -14,11 +14,11 @@ def setup_and_cleanup(): The part before the yield command is run before the test (setup) and the part after the yield command is run after (cleanup) """ - print("Setting up") + # Set up test_input_dir = INPUT_DIR test_output_dir = OUTPUT_DIR / "test_output" yield test_input_dir, test_output_dir - print("Cleaning up") + # Clean up shutil.rmtree(test_output_dir) diff --git a/tests/test_4_mappings.py b/tests/test_4_mappings.py index 84dd808..923565a 100644 --- a/tests/test_4_mappings.py +++ b/tests/test_4_mappings.py @@ -22,7 +22,7 @@ def setup_and_cleanup(): The part before the yield command is run before the test (setup) and the part after the yield command is run after (cleanup) """ - print("Setting up") + # Set up ledger_mapping['sample_bitcoin'] = DefaultMapping ledger_parser['sample_bitcoin'] = DefaultParser ledger_mapping['sample_ethereum'] = EthereumMapping @@ -35,7 +35,7 @@ def setup_and_cleanup(): test_input_dir = INPUT_DIR test_output_dir = OUTPUT_DIR / "test_output" yield pool_info_dir, test_input_dir, test_output_dir - print("Cleaning up") + # Clean up shutil.rmtree(test_output_dir) diff --git a/tests/test_5_analyze.py b/tests/test_5_analyze.py index 5807bc8..683549a 100644 --- a/tests/test_5_analyze.py +++ b/tests/test_5_analyze.py @@ -12,7 +12,7 @@ def setup_and_cleanup(): The part before the yield command is run before the test (setup) and the part after the yield command is run after (cleanup) """ - print("Setting up") + # Set up test_io_dir = OUTPUT_DIR / "test_output" test_bitcoin_dir = test_io_dir / "sample_bitcoin" if not os.path.exists(test_bitcoin_dir): @@ -36,7 +36,7 @@ def setup_and_cleanup(): with open(test_bitcoin_dir / f'{timeframe}.csv', 'w') as f: f.write(content) yield test_io_dir - print("Cleaning up") + # Clean up shutil.rmtree(test_io_dir) diff --git a/tests/test_6_end_to_end.py b/tests/test_6_end_to_end.py index bc4f824..c12044f 100644 --- a/tests/test_6_end_to_end.py +++ b/tests/test_6_end_to_end.py @@ -19,14 +19,14 @@ def setup_and_cleanup(): The part before the yield command is run before the test (setup) and the part after the yield command is run after (cleanup) """ - print("Setting up") + # Set up test_output_dir = OUTPUT_DIR / "test_output" ledger_mapping['sample_bitcoin'] = DefaultMapping ledger_parser['sample_bitcoin'] = DefaultParser ledger_mapping['sample_cardano'] = CardanoMapping ledger_parser['sample_cardano'] = DummyParser yield test_output_dir - print("Cleaning up") + # Clean up shutil.rmtree(test_output_dir)