[tor-commits] [sbws/master] v3bwfile: Include scanner country in the headers

[juga at torproject.org]

commit 6fc0810a5f9c0e54d20605bdb122a2b709ce6c19
Author: juga0 <juga at riseup.net>
Date:   Wed Feb 6 12:57:18 2019 +0000

    v3bwfile: Include scanner country in the headers
---
 sbws/core/generate.py           |  6 ++++--
 sbws/lib/v3bwfile.py            | 12 ++++++++----
 tests/unit/lib/test_v3bwfile.py |  7 +++++--
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/sbws/core/generate.py b/sbws/core/generate.py
index df86029..c022b5b 100644
--- a/sbws/core/generate.py
+++ b/sbws/core/generate.py
@@ -104,8 +104,10 @@ def main(args, conf):
     state_fpath = conf.getpath('paths', 'state_fname')
     consensus_path = os.path.join(conf.getpath('tor', 'datadir'),
                                   "cached-consensus")
-    bw_file = V3BWFile.from_results(results, state_fpath, args.scale_constant,
-                                    scaling_method,
+    # Accept None as scanner_country to be compatible with older versions.
+    scanner_country = conf['scanner'].get('country')
+    bw_file = V3BWFile.from_results(results, scanner_country, state_fpath,
+                                    args.scale_constant, scaling_method,
                                     torflow_cap=args.torflow_bw_margin,
                                     round_digs=args.round_digs,
                                     secs_recent=args.secs_recent,
diff --git a/sbws/lib/v3bwfile.py b/sbws/lib/v3bwfile.py
index 8cb9226..b0ab385 100644
--- a/sbws/lib/v3bwfile.py
+++ b/sbws/lib/v3bwfile.py
@@ -28,7 +28,8 @@ KEYVALUE_SEP_V1 = '='
 KEYVALUE_SEP_V2 = ' '
 # List of the extra KeyValues accepted by the class
 EXTRA_ARG_KEYVALUES = ['software', 'software_version', 'file_created',
-                       'earliest_bandwidth', 'generator_started']
+                       'earliest_bandwidth', 'generator_started',
+                       'scanner_country']
 STATS_KEYVALUES = ['number_eligible_relays', 'minimum_number_eligible_relays',
                    'number_consensus_relays', 'percent_eligible_relays',
                    'minimum_percent_eligible_relays']
@@ -140,7 +141,7 @@ class V3BWHeader(object):
         return self.strv2
 
     @classmethod
-    def from_results(cls, results, state_fpath=''):
+    def from_results(cls, results, scanner_country=None, state_fpath=''):
         kwargs = dict()
         latest_bandwidth = cls.latest_bandwidth_from_results(results)
         earliest_bandwidth = cls.earliest_bandwidth_from_results(results)
@@ -150,6 +151,9 @@ class V3BWHeader(object):
         kwargs['earliest_bandwidth'] = unixts_to_isodt_str(earliest_bandwidth)
         if generator_started is not None:
             kwargs['generator_started'] = generator_started
+        # To be compatible with older bandwidth files, do not require it.
+        if scanner_country is not None:
+            kwargs['scanner_country'] = scanner_country
         h = cls(timestamp, **kwargs)
         return h
 
@@ -524,7 +528,7 @@ class V3BWFile(object):
                                            for bw_line in self.bw_lines])
 
     @classmethod
-    def from_results(cls, results, state_fpath='',
+    def from_results(cls, results, scanner_country=None, state_fpath='',
                      scale_constant=SBWS_SCALE_CONSTANT,
                      scaling_method=TORFLOW_SCALING,
                      torflow_obs=TORFLOW_OBS_LAST,
@@ -551,7 +555,7 @@ class V3BWFile(object):
 
         """
         log.info('Processing results to generate a bandwidth list file.')
-        header = V3BWHeader.from_results(results, state_fpath)
+        header = V3BWHeader.from_results(results, scanner_country, state_fpath)
         bw_lines_raw = []
         number_consensus_relays = cls.read_number_consensus_relays(
             consensus_path)
diff --git a/tests/unit/lib/test_v3bwfile.py b/tests/unit/lib/test_v3bwfile.py
index 2d38f48..506f63b 100644
--- a/tests/unit/lib/test_v3bwfile.py
+++ b/tests/unit/lib/test_v3bwfile.py
@@ -16,6 +16,8 @@ from sbws.util.timestamp import now_fname, now_isodt_str, now_unixts
 timestamp = 1523974147
 timestamp_l = str(timestamp)
 version_l = KEYVALUE_SEP_V1.join(['version', SPEC_VERSION])
+scanner_country = 'US'
+scanner_country_l = KEYVALUE_SEP_V1.join(['scanner_country', scanner_country])
 software_l = KEYVALUE_SEP_V1.join(['software', 'sbws'])
 software_version_l = KEYVALUE_SEP_V1.join(['software_version', version])
 file_created = '2018-04-25T13:10:57'
@@ -24,7 +26,7 @@ latest_bandwidth = '2018-04-17T14:09:07'
 latest_bandwidth_l = KEYVALUE_SEP_V1.join(['latest_bandwidth',
                                           latest_bandwidth])
 header_ls = [timestamp_l, version_l, file_created_l, latest_bandwidth_l,
-             software_l, software_version_l, TERMINATOR]
+             scanner_country_l, software_l, software_version_l, TERMINATOR]
 header_str = LINE_SEP.join(header_ls) + LINE_SEP
 earliest_bandwidth = '2018-04-16T14:09:07'
 earliest_bandwidth_l = KEYVALUE_SEP_V1.join(['earliest_bandwidth',
@@ -53,7 +55,8 @@ v3bw_str = header_extra_str + raw_bwl_str
 
 def test_v3bwheader_str():
     """Test header str"""
-    header = V3BWHeader(timestamp_l, file_created=file_created)
+    header = V3BWHeader(timestamp_l, scanner_country=scanner_country,
+                        file_created=file_created)
     assert header_str == str(header)