[tor-commits] [ooni-probe/master] Make it run once if there are no Asset

[art at torproject.org]

commit a925c465b45907d675ca915e3b5c2832b39390e8
Author: Arturo Filastò <hellais at torproject.org>
Date:   Thu May 31 21:15:24 2012 +0200

    Make it run once if there are no Asset
    * Write regression test for it
---
 ooni/ooniprobe.py      |    6 +++---
 ooni/plugoo/reports.py |   43 ++++++++++++++++++++++++++++++-------------
 ooni/plugoo/work.py    |   10 ++++++++--
 tests/test_worker.py   |   11 +++++++++++
 4 files changed, 52 insertions(+), 18 deletions(-)

diff --git a/ooni/ooniprobe.py b/ooni/ooniprobe.py
index d75b9b9..7248e70 100755
--- a/ooni/ooniprobe.py
+++ b/ooni/ooniprobe.py
@@ -59,12 +59,12 @@ def runTest(test, options, global_options):
 
     parallelism = int(global_options['parallelism'])
     worker = work.Worker(parallelism)
-    test = plugoo[test].__class__
-    report = reports.Report(global_options['output'])
+    test_class = plugoo[test].__class__
+    report = reports.Report(test, global_options['output'])
 
     log.start(global_options['log'], 1)
 
-    wgen = work.WorkGenerator(test(options, global_options, report),
+    wgen = work.WorkGenerator(test_class(options, global_options, report),
                               dict(options),
                               start=options['resume'])
 
diff --git a/ooni/plugoo/reports.py b/ooni/plugoo/reports.py
index 099aff7..6140294 100644
--- a/ooni/plugoo/reports.py
+++ b/ooni/plugoo/reports.py
@@ -1,9 +1,8 @@
 import os
-from datetime import datetime
 import yaml
 
 import itertools
-from ooni import log
+from ooni import log, date
 
 class Report:
     """This is the ooni-probe reporting mechanism. It allows
@@ -19,10 +18,11 @@ class Report:
          inbound connection and accept a stream of data (think of it
          as a `nc -l -p <port> > filename.txt`)
     """
-    def __init__(self, file="test.report",
+    def __init__(self, testname=None, file="report.log",
                  scp=None,
                  tcp=None):
 
+        self.testname = testname
         self.file = file
         self.tcp = tcp
         self.scp = scp
@@ -36,16 +36,21 @@ class Report:
         #    print self.file
 
         self.scp = None
-
-    def __call__(self, data):
-        """
-        This should be invoked every time you wish to write some
-        data to the reporting system
-        """
-        #print "Writing report(s)"
-        #dump = '--- \n'
-        dump = yaml.dump([data])
-        #dump += yaml.dump(data)
+        self.write_header()
+
+    def write_header(self):
+        pretty_date = date.pretty_date()
+        header = "# OONI Probe Report for Test %s\n" % self.testname
+        header += "# %s\n\n" % pretty_date
+        self._write_to_report(header)
+        # XXX replace this with something proper
+        test_details = {'start_time': date.now(),
+                        'asn': 'ASN-1234',
+                        'test_name': self.testname,
+                        'addr': '1234'}
+        self(test_details)
+
+    def _write_to_report(self, dump):
         reports = []
 
         if self.file:
@@ -61,6 +66,18 @@ class Report:
         for report in reports:
             self.send_report(dump, report)
 
+    def __call__(self, data):
+        """
+        This should be invoked every time you wish to write some
+        data to the reporting system
+        """
+        #print "Writing report(s)"
+        #dump = '--- \n'
+        dump = yaml.dump([data])
+        #dump += yaml.dump(data)
+
+        self._write_to_report(dump)
+
     def file_report(self, data, file=None, mode='a+'):
         """
         This reports to a file in YAML format
diff --git a/ooni/plugoo/work.py b/ooni/plugoo/work.py
index b9fc6b8..8ba0a3f 100644
--- a/ooni/plugoo/work.py
+++ b/ooni/plugoo/work.py
@@ -130,7 +130,11 @@ class WorkGenerator(object):
     def __init__(self, test, arguments=None, start=None):
         self.Test = test
 
-        self.assetGenerator = itertools.product(*self.Test.assets.values())
+        if self.Test.assets and self.Test.assets.values()[0]:
+            self.assetGenerator = itertools.product(*self.Test.assets.values())
+        else:
+            self.assetGenerator = None
+
         self.assetNames = self.Test.assets.keys()
 
         self.idx = 0
@@ -148,6 +152,9 @@ class WorkGenerator(object):
             self.idx += 1
 
     def next(self):
+        if self.end:
+            raise StopIteration
+
         if not self.assetGenerator:
             self.end = True
             return (self.assetNames, self.Test, self.idx)
@@ -161,7 +168,6 @@ class WorkGenerator(object):
             raise StopIteration
 
         self.idx += 1
-        print "IDX:%s" % self.idx
         return (ret, self.Test, self.idx)
 
     def p_next(self):
diff --git a/tests/test_worker.py b/tests/test_worker.py
index e53750b..3222f65 100644
--- a/tests/test_worker.py
+++ b/tests/test_worker.py
@@ -12,7 +12,18 @@ class WorkerTestCase(unittest.TestCase):
             asset.append(i)
         dummytest.assets['asset'] = asset
         wgen = work.WorkGenerator(dummytest)
+
         for j, x in enumerate(wgen):
             pass
         self.assertEqual(i, j)
 
+    def testNoAssets(self):
+        class DummyTest:
+            assets = {'asset': None}
+        dummytest = DummyTest()
+        wgen = work.WorkGenerator(dummytest)
+        i = 0
+        for j in wgen:
+            i += 1
+        self.assertEqual(i, 1)
+