Fix issue and add tests (#54)

* Record exclude_below_fees flag in output

* Test exclude_below_fees in get_output_row

* Test Theil with empty input

* Test run

* Fix bugs in test analyze

* Test write_csv_output

* Remove unused line

* Add map tests

* Test apply_mapping

* Test fill_db_with_balances

* Minor change for code clarity

Author: dimkarakostas

tests/test_analyze.py

@@ -1,5 +1,6 @@
-from tokenomics_decentralization.analyze import get_output_row, analyze_snapshot, analyze
+from tokenomics_decentralization.analyze import get_output_row, analyze_snapshot, analyze, write_csv_output
 import pathlib
+import os
 
 
 def test_get_output_row(mocker):
@@ -20,22 +21,28 @@ def test_get_output_row(mocker):
 
     metrics = {'hhi': 1, 'gini': 0}
     csv_row = get_output_row('bitcoin', '2010-01-01', metrics)
-    assert csv_row == ['bitcoin', '2010-01-01', False, False, 'absolute', 0, 1, 0]
+    assert csv_row == ['bitcoin', '2010-01-01', False, False, 'absolute', 0, False, 1, 0]
 
     get_no_clustering_mock.return_value = True
     metrics = {'non-clustered hhi': 1, 'non-clustered gini': 0}
     csv_row = get_output_row('bitcoin', '2010-01-01', metrics)
-    assert csv_row == ['bitcoin', '2010-01-01', True, False, 'absolute', 0, 1, 0]
+    assert csv_row == ['bitcoin', '2010-01-01', True, False, 'absolute', 0, False, 1, 0]
 
     get_exclude_contracts_mock.return_value = True
     metrics = {'exclude_contracts non-clustered hhi': 1, 'exclude_contracts non-clustered gini': 0}
     csv_row = get_output_row('bitcoin', '2010-01-01', metrics)
-    assert csv_row == ['bitcoin', '2010-01-01', True, True, 'absolute', 0, 1, 0]
+    assert csv_row == ['bitcoin', '2010-01-01', True, True, 'absolute', 0, False, 1, 0]
 
     get_top_limit_value_mock.return_value = 1
     metrics = {'top-1_absolute exclude_contracts non-clustered hhi': 1, 'top-1_absolute exclude_contracts non-clustered gini': 0}
     csv_row = get_output_row('bitcoin', '2010-01-01', metrics)
-    assert csv_row == ['bitcoin', '2010-01-01', True, True, 'absolute', 1, 1, 0]
+    assert csv_row == ['bitcoin', '2010-01-01', True, True, 'absolute', 1, False, 1, 0]
+
+    get_exclude_below_fees_mock.return_value = True
+    get_top_limit_value_mock.return_value = 1
+    metrics = {'top-1_absolute exclude_below_fees exclude_contracts non-clustered hhi': 1, 'top-1_absolute exclude_below_fees exclude_contracts non-clustered gini': 0}
+    csv_row = get_output_row('bitcoin', '2010-01-01', metrics)
+    assert csv_row == ['bitcoin', '2010-01-01', True, True, 'absolute', 1, True, 1, 0]
 
 
 def test_analyze_snapshot(mocker):
@@ -54,6 +61,9 @@ def test_analyze_snapshot(mocker):
     db_insert_mock = mocker.patch('tokenomics_decentralization.db_helper.insert_metric')
     db_commit_mock = mocker.patch('tokenomics_decentralization.db_helper.commit_database')
 
+    compute_hhi_mock = mocker.patch('tokenomics_decentralization.analyze.compute_hhi')
+    compute_tau_mock = mocker.patch('tokenomics_decentralization.analyze.compute_tau')
+
     get_force_analyze_mock.return_value = False
     get_no_clustering_mock.return_value = False
     get_exclude_contracts_mock.return_value = False
@@ -88,41 +98,94 @@ def test_analyze_snapshot(mocker):
     db_insert_mock.return_value = None
     db_commit_mock.return_value = None
 
+    compute_hhi_mock.return_value = 2
     output = analyze_snapshot(None, 'bitcoin', '2010-01-01')
-    assert output == {'top-1_absolute exclude_below_fees exclude_contracts non-clustered hhi': 10000}
+    assert output == {'top-1_absolute exclude_below_fees exclude_contracts non-clustered hhi': 2}
 
     get_no_clustering_mock.return_value = False
 
+    compute_hhi_mock.return_value = 3
     output = analyze_snapshot(None, 'bitcoin', '2010-01-01')
-    assert output == {'top-1_absolute exclude_below_fees exclude_contracts hhi': 10000}
+    assert output == {'top-1_absolute exclude_below_fees exclude_contracts hhi': 3}
 
     get_top_limit_value_mock.return_value = 0
 
+    compute_hhi_mock.return_value = 4
     output = analyze_snapshot(None, 'bitcoin', '2010-01-01')
-    assert output == {'exclude_below_fees exclude_contracts hhi': 5000}
+    assert output == {'exclude_below_fees exclude_contracts hhi': 4}
 
     get_top_limit_type_mock.return_value = 'percentage'
     get_top_limit_value_mock.return_value = 0.5
 
+    compute_hhi_mock.return_value = 5
     output = analyze_snapshot(None, 'bitcoin', '2010-01-01')
-    assert output == {'top-0.5_percentage exclude_below_fees exclude_contracts hhi': 10000}
+    assert output == {'top-0.5_percentage exclude_below_fees exclude_contracts hhi': 5}
 
     get_top_limit_value_mock.return_value = 0
     get_metrics_mock.return_value = ['tau=0.5']
 
+    compute_tau_mock.return_value = [100, None]
     output = analyze_snapshot(None, 'bitcoin', '2010-01-01')
-    assert output == {'exclude_below_fees exclude_contracts tau=0.5': 1}
+    assert output == {'exclude_below_fees exclude_contracts tau=0.5': 100}
+
+
+def test_write_csv_output(mocker):
+    get_metrics_mock = mocker.patch('tokenomics_decentralization.helper.get_metrics')
+    get_metrics_mock.return_value = ['hhi']
+
+    get_output_directories_mock = mocker.patch('tokenomics_decentralization.helper.get_output_directories')
+    get_output_directories_mock.return_value = [pathlib.Path(__file__).resolve().parent]
+
+    get_no_clustering_mock = mocker.patch('tokenomics_decentralization.helper.get_no_clustering_flag')
+    get_exclude_contracts_mock = mocker.patch('tokenomics_decentralization.helper.get_exclude_contracts_flag')
+    get_exclude_below_fees_mock = mocker.patch('tokenomics_decentralization.helper.get_exclude_below_fees_flag')
+    get_top_limit_type_mock = mocker.patch('tokenomics_decentralization.helper.get_top_limit_type')
+    get_top_limit_value_mock = mocker.patch('tokenomics_decentralization.helper.get_top_limit_value')
+
+    get_no_clustering_mock.return_value = False
+    get_exclude_contracts_mock.return_value = False
+    get_exclude_below_fees_mock.return_value = False
+    get_top_limit_type_mock.return_value = 'absolute'
+    get_top_limit_value_mock.return_value = 0
+
+    write_csv_output([
+        ['bitcoin', '2010-01-01', False, False, 'absolute', 0, False, 100],
+        ['ethereum', '2010-01-01', False, False, 'absolute', 0, False, 200],
+        ])
+    with open(pathlib.Path(__file__).resolve().parent / 'output.csv') as f:
+        lines = f.readlines()
+        assert lines[0] == ','.join(['ledger', 'snapshot date', 'no_clustering', 'exclude_contract_addresses', 'top_limit_type', 'top_limit_value', 'exclude_below_fees', 'hhi']) + '\n'
+        assert lines[1] == ','.join(['bitcoin', '2010-01-01', 'False', 'False', 'absolute', '0', 'False', '100']) + '\n'
+        assert lines[2] == ','.join(['ethereum', '2010-01-01', 'False', 'False', 'absolute', '0', 'False', '200']) + '\n'
+    os.remove(pathlib.Path(__file__).resolve().parent / 'output.csv')
+
+    get_no_clustering_mock.return_value = True
+    get_exclude_contracts_mock.return_value = True
+    get_exclude_below_fees_mock.return_value = True
+    get_top_limit_type_mock.return_value = 'absolute'
+    get_top_limit_value_mock.return_value = 10
+
+    write_csv_output([
+        ['bitcoin', '2010-01-01', False, False, 'absolute', 0, False, 100],
+        ['ethereum', '2010-01-01', False, False, 'absolute', 0, False, 200],
+        ])
+    with open(pathlib.Path(__file__).resolve().parent / 'output-no_clustering-exclude_contract_addresses-absolute_10-exclude_below_fees.csv') as f:
+        lines = f.readlines()
+        assert lines[0] == ','.join(['ledger', 'snapshot date', 'no_clustering', 'exclude_contract_addresses', 'top_limit_type', 'top_limit_value', 'exclude_below_fees', 'hhi']) + '\n'
+        assert lines[1] == ','.join(['bitcoin', '2010-01-01', 'False', 'False', 'absolute', '0', 'False', '100']) + '\n'
+        assert lines[2] == ','.join(['ethereum', '2010-01-01', 'False', 'False', 'absolute', '0', 'False', '200']) + '\n'
+    os.remove(pathlib.Path(__file__).resolve().parent / 'output-no_clustering-exclude_contract_addresses-absolute_10-exclude_below_fees.csv')
 
 
 def test_analyze(mocker):
     get_output_directories_mock = mocker.patch('tokenomics_decentralization.helper.get_output_directories')
-    get_output_directories_mock.return_value = [pathlib.Path(__file__).resolve()]
+    get_output_directories_mock.return_value = [pathlib.Path(__file__).resolve().parent]
 
     is_file_mock = mocker.patch('os.path.isfile')
     is_file_mock.return_value = True
 
-    get_db_connector_mock = mocker.patch('tokenomics_decentralization.schema.get_connector')
-    get_db_connector_mock.return_value = None
+    get_db_connector_mock = mocker.patch('tokenomics_decentralization.analyze.get_connector')
+    get_db_connector_mock.return_value = 'connector'
 
     analyze_snapshot_mock = mocker.patch('tokenomics_decentralization.analyze.analyze_snapshot')
     analyze_snapshot_mock.return_value = {'hhi': 1}
@@ -135,7 +198,7 @@ def test_analyze(mocker):
 
     output_rows = analyze(['bitcoin'], ['2010-01-01'])
     assert output_rows == ['row']
-    analyze_snapshot_mock.assert_called_with(None, 'bitcoin', '2010-01-01')
+    analyze_snapshot_mock.assert_called_with('connector', 'bitcoin', '2010-01-01')
 
     output_rows = analyze(['bitcoin', 'ethereum'], ['2010-01-01', '2011-01-01'])
     assert output_rows == ['row', 'row', 'row', 'row']

tests/test_map.py

@@ -0,0 +1,138 @@
+from tokenomics_decentralization.map import fill_db_with_addresses, apply_mapping, fill_db_with_balances
+from unittest.mock import call
+import pathlib
+
+
+def test_fill_db_with_addresses(mocker):
+    db_insert_ledger_mock = mocker.patch('tokenomics_decentralization.db_helper.insert_ledger')
+    db_insert_ledger_mock.return_value = None
+    db_insert_entity_mock = mocker.patch('tokenomics_decentralization.db_helper.insert_entity')
+    db_insert_entity_mock.return_value = None
+    db_insert_update_address_mock = mocker.patch('tokenomics_decentralization.db_helper.insert_update_address')
+    db_insert_update_address_mock.return_value = None
+    db_commit_mock = mocker.patch('tokenomics_decentralization.db_helper.commit_database')
+    db_commit_mock.return_value = None
+
+    json_load_mock = mocker.patch('json.load')
+
+    json_load_mock.return_value = {
+        'address 1': {'name': 'Entity Name 1', 'source': 'website'},
+        'address 2': {'name': 'Entity Name 2', 'source': 'website', 'is_contract': True},
+    }
+
+    fill_db_with_addresses('connector', 'random ledger')
+    assert json_load_mock.call_args_list == []
+    assert db_insert_entity_mock.call_args_list == []
+    assert db_insert_update_address_mock.call_args_list == []
+    assert db_commit_mock.call_args_list == []
+
+    fill_db_with_addresses('connector', 'bitcoin')
+    db_insert_ledger_mock.assert_called_with('connector', 'bitcoin')
+    assert call('connector', 'bitcoin', 'Entity Name 1') in db_insert_entity_mock.call_args_list
+    assert call('connector', 'bitcoin', 'Entity Name 2') in db_insert_entity_mock.call_args_list
+    assert call('connector', 'bitcoin', 'address 1', 'Entity Name 1', False) in db_insert_update_address_mock.call_args_list
+    assert call('connector', 'bitcoin', 'address 2', 'Entity Name 2', True) in db_insert_update_address_mock.call_args_list
+    db_commit_mock.assert_called()
+
+
+def test_fill_db_with_balances(mocker):
+    get_input_directories_mock = mocker.patch('tokenomics_decentralization.helper.get_input_directories')
+    get_input_directories_mock.return_value = [pathlib.Path('/input').resolve()]
+
+    os_isfile_mock = mocker.patch('os.path.isfile')
+    os_isfile_mock.return_value = True
+
+    db_insert_address_without_update_mock = mocker.patch('tokenomics_decentralization.db_helper.insert_address_without_update')
+    db_insert_snapshot_mock = mocker.patch('tokenomics_decentralization.db_helper.insert_snapshot')
+    db_insert_balance_mock = mocker.patch('tokenomics_decentralization.db_helper.insert_balance')
+    db_update_circulation_mock = mocker.patch('tokenomics_decentralization.db_helper.update_circulation')
+    db_commit_mock = mocker.patch('tokenomics_decentralization.db_helper.commit_database')
+
+    mocker.patch('builtins.open', mocker.mock_open(read_data='address,type,balance\naddr1,pubkey,100\naddr2,pubkey,200'))
+
+    fill_db_with_balances('connector', 'bitcoin', '2010-01-01')
+    assert db_insert_snapshot_mock.call_args_list == [call('connector', 'bitcoin', '2010-01-01')]
+    assert db_insert_address_without_update_mock.call_args_list == [call('connector', 'bitcoin', 'addr1'), call('connector', 'bitcoin', 'addr2')]
+    assert db_insert_balance_mock.call_args_list == [call('connector', 'bitcoin', '2010-01-01', 'addr1', 100), call('connector', 'bitcoin', '2010-01-01', 'addr2', 200)]
+    assert db_update_circulation_mock.call_args_list == [call('connector', 'bitcoin', '2010-01-01', 300)]
+    assert db_commit_mock.call_args_list == [call()]
+
+    db_insert_address_without_update_mock.reset_mock()
+    db_insert_snapshot_mock.reset_mock()
+    db_insert_balance_mock.reset_mock()
+    db_update_circulation_mock.reset_mock()
+    db_commit_mock.reset_mock()
+
+    mocker.patch('builtins.open', mocker.mock_open(read_data='address,balance\naddr1,0\naddr2,1000000000'))
+
+    fill_db_with_balances('connector', 'ethereum', '2010-01-01')
+    assert db_insert_snapshot_mock.call_args_list == [call('connector', 'ethereum', '2010-01-01')]
+    assert db_insert_address_without_update_mock.call_args_list == [call('connector', 'ethereum', 'addr2')]
+    assert db_insert_balance_mock.call_args_list == [call('connector', 'ethereum', '2010-01-01', 'addr2', 1)]
+    assert db_update_circulation_mock.call_args_list == [call('connector', 'ethereum', '2010-01-01', 1)]
+    assert db_commit_mock.call_args_list == [call()]
+
+
+def test_apply_mapping(mocker):
+    get_db_connector_mock = mocker.patch('tokenomics_decentralization.map.get_connector')
+    get_db_connector_mock.return_value = 'connector'
+
+    fill_db_with_addresses_mock = mocker.patch('tokenomics_decentralization.map.fill_db_with_addresses')
+    fill_db_with_balances_mock = mocker.patch('tokenomics_decentralization.map.fill_db_with_balances')
+
+    get_force_map_addresses_mock = mocker.patch('tokenomics_decentralization.helper.get_force_map_addresses_flag')
+    get_force_map_balances_mock = mocker.patch('tokenomics_decentralization.helper.get_force_map_balances_flag')
+
+    get_output_directories_mock = mocker.patch('tokenomics_decentralization.helper.get_output_directories')
+    get_input_directories_mock = mocker.patch('tokenomics_decentralization.helper.get_input_directories')
+
+    os_isfile_mock = mocker.patch('os.path.isfile')
+    os_isfile_mock.side_effect = {
+        pathlib.Path('/output').resolve() / 'bitcoin_2010-01-01.db': False,
+        pathlib.Path('/input1').resolve() / 'bitcoin_2010-01-01_raw_data.csv': False,
+        pathlib.Path('/input2').resolve() / 'bitcoin_2010-01-01_raw_data.csv': False,
+    }.get
+
+    get_force_map_addresses_mock.return_value = False
+    get_force_map_balances_mock.return_value = False
+
+    get_output_directories_mock.return_value = [pathlib.Path('/output').resolve()]
+    get_input_directories_mock.return_value = [pathlib.Path('/input1').resolve(), pathlib.Path('/input2').resolve()]
+
+    apply_mapping('bitcoin', '2010-01-01')
+    assert get_db_connector_mock.call_args_list == []
+    assert fill_db_with_addresses_mock.call_args_list == []
+    assert fill_db_with_balances_mock.call_args_list == []
+
+    os_isfile_mock.side_effect = {
+        pathlib.Path('/output').resolve() / 'bitcoin_2010-01-01.db': True,
+    }.get
+    get_input_directories_mock.reset_mock()
+
+    apply_mapping('bitcoin', '2010-01-01')
+    assert get_input_directories_mock.call_args_list == []
+    assert fill_db_with_addresses_mock.call_args_list == []
+    assert fill_db_with_balances_mock.call_args_list == []
+
+    get_force_map_addresses_mock.return_value = True
+    get_force_map_balances_mock.return_value = True
+
+    apply_mapping('bitcoin', '2010-01-01')
+    assert fill_db_with_addresses_mock.call_args_list == [call('connector', 'bitcoin')]
+    assert fill_db_with_balances_mock.call_args_list == [call('connector', 'bitcoin', '2010-01-01')]
+
+    os_isfile_mock.side_effect = {
+        pathlib.Path('/output').resolve() / 'bitcoin_2010-01-01.db': False,
+        pathlib.Path('/input1').resolve() / 'bitcoin_2010-01-01_raw_data.csv': True
+    }.get
+
+    get_force_map_addresses_mock.return_value = False
+    get_force_map_balances_mock.return_value = False
+    get_input_directories_mock.reset_mock()
+    fill_db_with_addresses_mock.reset_mock()
+    fill_db_with_balances_mock.reset_mock()
+
+    apply_mapping('bitcoin', '2010-01-01')
+    assert get_input_directories_mock.call_args_list == [call()]
+    assert fill_db_with_addresses_mock.call_args_list == [call('connector', 'bitcoin')]
+    assert fill_db_with_balances_mock.call_args_list == [call('connector', 'bitcoin', '2010-01-01')]

tests/test_metrics.py

@@ -173,3 +173,7 @@ def test_compute_theil_index():
     tokens_per_entity = {('a', 432)}
     theil_t = compute_theil_index(tokens_per_entity, 432)
     assert round(theil_t, decimals) == 0
+
+    tokens_per_entity = {}
+    theil_t = compute_theil_index(tokens_per_entity, 432)
+    assert theil_t == 0

tests/test_run.py

@@ -0,0 +1,27 @@
+import run
+from unittest.mock import call
+
+
+def test_run(mocker):
+    apply_mapping_mock = mocker.patch('run.apply_mapping')
+    analyze_mock = mocker.patch('run.analyze')
+    plot_mock = mocker.patch('run.plot')
+    get_plot_flag_mock = mocker.patch('tokenomics_decentralization.helper.get_plot_flag')
+
+    get_plot_flag_mock.return_value = False
+
+    run.main(['bitcoin'], ['2010-01-01'])
+
+    apply_mapping_mock.assert_called_with('bitcoin', '2010-01-01')
+    analyze_mock.assert_called_with(['bitcoin'], ['2010-01-01'])
+
+    run.main(['bitcoin', 'ethereum'], ['2010-01-01', '2020-01-01'])
+    assert call('bitcoin', '2010-01-01') in apply_mapping_mock.call_args_list
+    assert call('bitcoin', '2020-01-01') in apply_mapping_mock.call_args_list
+    assert call('ethereum', '2010-01-01') in apply_mapping_mock.call_args_list
+    assert call('ethereum', '2020-01-01') in apply_mapping_mock.call_args_list
+    analyze_mock.assert_called_with(['bitcoin', 'ethereum'], ['2010-01-01', '2020-01-01'])
+
+    get_plot_flag_mock.return_value = True
+    run.main(['bitcoin'], ['2010-01-01'])
+    plot_mock.assert_called()

tokenomics_decentralization/analyze.py

@@ -95,7 +95,7 @@ def get_output_row(ledger, date, metrics):
     top_limit_type = hlp.get_top_limit_type()
     top_limit_value = hlp.get_top_limit_value()
 
-    csv_row = [ledger, date, no_clustering, exclude_contract_addresses_flag, top_limit_type, top_limit_value]
+    csv_row = [ledger, date, no_clustering, exclude_contract_addresses_flag, top_limit_type, top_limit_value, exclude_below_fees_flag]
 
     for metric_name in hlp.get_metrics():
         val = metric_name
@@ -112,20 +112,23 @@ def get_output_row(ledger, date, metrics):
 
 
 def write_csv_output(output_rows):
-    header = ['ledger', 'snapshot date', 'no_clustering', 'exclude_contract_addresses', 'top_limit_type', 'top_limit_value']
+    header = ['ledger', 'snapshot date', 'no_clustering', 'exclude_contract_addresses', 'top_limit_type', 'top_limit_value', 'exclude_below_fees']
     header += hlp.get_metrics()
 
     no_clustering = hlp.get_no_clustering_flag()
     exclude_contract_addresses_flag = hlp.get_exclude_contracts_flag()
     top_limit_type = hlp.get_top_limit_type()
     top_limit_value = hlp.get_top_limit_value()
+    exclude_below_fees_flag = hlp.get_exclude_below_fees_flag()
     output_filename = 'output'
     if no_clustering:
         output_filename += '-no_clustering'
     if exclude_contract_addresses_flag:
         output_filename += '-exclude_contract_addresses'
     if top_limit_value:
         output_filename += f'-{top_limit_type}_{top_limit_value}'
+    if exclude_below_fees_flag:
+        output_filename += '-exclude_below_fees'
     output_filename += '.csv'
 
     output_dir = hlp.get_output_directories()[0]

tokenomics_decentralization/map.py

@@ -11,7 +11,6 @@
 
 def fill_db_with_addresses(conn, ledger):
     db_hlp.insert_ledger(conn, ledger)
-    db_hlp.get_ledger_id(conn, ledger)
 
     try:
         with open(hlp.MAPPING_INFO_DIR / f'addresses/{ledger}.json') as f: