#!/usr/bin/python

import argparse
import csv
import subprocess
import sys
import time
import urllib

import StringIO

FLOOD_SIZE_MB = 20
FLOOD_MAX_SPEED = '5M'

parser = argparse.ArgumentParser(description='Measure SOCKS proxy throughput across browser versions.')
parser.add_argument('clone_path', help='path to pre-built uproxy-lib repo')
args = parser.parse_args()

# Where is flood server?
flood_ip = subprocess.check_output('./flood.sh ' + str(FLOOD_SIZE_MB) + ' ' + FLOOD_MAX_SPEED, shell=True).strip()
print '** using flood server at ' + flood_ip

browsers = ['chrome', 'firefox']
versions = ['stable', 'beta', 'canary']

# Run the benchmarks.
throughput = {}
for browser in browsers:
  for version in versions:
    print '** ' + browser + '/' + version

    # Start the relevant config, stopping the previous one if necessary.
    # TODO: check first if running, to avoid spurious warnings
    subprocess.call(['docker', 'rm', '-f', 'uproxy-getter', 'uproxy-giver'])
    # TODO: use python 3.3+ timeouts
    spec = browser + '-' + version
    subprocess.call('./run_pair.sh -p ' + args.clone_path + ' ' + spec + ' ' + spec, shell=True)

    # time.time is good for Unix-like systems.
    start = time.time()
    subprocess.call(['nc', '-X', '5', '-x', 'localhost:9999', flood_ip, '1224'])
    end = time.time()

    elapsed = round(end - start, 2)
    throughput[browser, version] = int((FLOOD_SIZE_MB / elapsed) * 1000)

    print '** throughput for ' + browser + '/' + version + ': ' + str(throughput[browser, version]) + 'K/sec'

# Raw summary.
print '** raw numbers: ' + str(throughput)

# CSV, e.g.:
#   throughput,stable,beta,canary
#   chrome,200,250,400
#   firefox,500,500,600
stringio = StringIO.StringIO()
writer = csv.writer(stringio)
headers = ['throughput']
headers.extend(versions)
writer.writerow(headers)
for browser in browsers:
  row = [browser]
  for version in versions:
    row.append(throughput[browser, version])
  writer.writerow(row)
print '** CSV'
print stringio.getvalue()

# URL which uses Datacopia's oh-so-simple GET-based API:
print '** http://www.datacopia.com/?data=' + urllib.quote(stringio.getvalue())

@@ -10,16 +10,17 @@
 n/python
+3
 %0A%0Aimport
@@ -40,16 +40,26 @@
 ort csv%0A
+import io%0A
 import s
@@ -108,25 +108,14 @@
 llib
-%0A%0Aimport StringIO
+.parse
 %0A%0AFL
@@ -493,16 +493,41 @@
 ell=True
+, universal_newlines=True
 ).strip(
@@ -529,25 +529,25 @@
 trip()%0Aprint
- 
+(
 '** using fl
@@ -564,24 +564,28 @@
  at ' + 
+str(
 flood_ip
 %0A%0Abrowse
@@ -576,16 +576,18 @@
 flood_ip
+))
 %0A%0Abrowse
@@ -749,25 +749,25 @@
 s:%0A    print
- 
+(
 '** ' + brow
@@ -785,18 +785,45 @@
  version
+)
 %0A%0A
+    result = 0%0A    try:%0A  
     # St
@@ -887,24 +887,26 @@
 essary.%0A    
+  
 # TODO: chec
@@ -952,16 +952,18 @@
 arnings%0A
+  
     subp
@@ -1037,43 +1037,8 @@
 )%0A  
-  # TODO: use python 3.3+ timeouts%0A
     
@@ -1068,16 +1068,18 @@
 version%0A
+  
     subp
@@ -1165,19 +1165,33 @@
 ell=True
+, timeout=15
 )%0A%0A
+  
     # ti
@@ -1229,16 +1229,18 @@
 ystems.%0A
+  
     star
@@ -1255,24 +1255,29 @@
 .time()%0A    
+  if 
 subprocess.c
@@ -1340,17 +1340,92 @@
 '1224'%5D)
-%0A
+ != 0:%0A        raise Exception('nc failed, proxy probably did not start')%0A  
     end 
@@ -1439,16 +1439,18 @@
 time()%0A%0A
+  
     elap
@@ -1485,36 +1485,16 @@
     
-throughput%5Bbrowser, version%5D
+  result
  = i
@@ -1535,22 +1535,24 @@
 0)%0A%0A    
+  
 print
- 
+(
 '** thro
@@ -1595,32 +1595,159 @@
 on + ': ' + str(
+result) + 'K/sec')%0A    except Exception as e:%0A      print('** failed to test ' + browser + '/' + version + ': ' + str(e))%0A%0A    
 throughput%5Bbrows
@@ -1758,27 +1758,25 @@
 version%5D
-) + 'K/sec'
+ = result
 %0A%0A# Raw 
@@ -1785,25 +1785,25 @@
 mmary.%0Aprint
- 
+(
 '** raw numb
@@ -1826,16 +1826,17 @@
 oughput)
+)
 %0A%0A# CSV,
@@ -1934,24 +1934,18 @@
 ingio = 
-StringIO
+io
 .StringI
@@ -2200,17 +2200,17 @@
 rint
- 
+(
 '** CSV'
 %0Apri
@@ -2205,23 +2205,24 @@
 '** CSV'
+)
 %0Aprint
- 
+(
 stringio
@@ -2232,16 +2232,17 @@
 tvalue()
+)
 %0A%0A# URL 
@@ -2297,17 +2297,17 @@
 I:%0Aprint
- 
+(
 '** http
@@ -2344,16 +2344,22 @@
  urllib.
+parse.
 quote(st
@@ -2356,29 +2356,30 @@
 e.quote(stringio.getvalue())
+)
 %0A

switch to python3 for timeout in subprocess