TestON/tests/HA/HAscaling/dependencies/onos-gen-partitions

#!/usr/bin/env python
"""
  Generate the partitions json file from the $OC* environment variables

  Usage: onos-gen-partitions output_file [num_nodes] [-e]
  If output file is not provided, the json is written to stdout.
"""

from os import environ
from collections import deque, OrderedDict
import re
import json
import sys
import hashlib

convert = lambda text: int(text) if text.isdigit() else text.lower()
alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]

def get_OC_vars():
  vars = []
  for var in environ:
    if re.match(r"OC[0-9]+", var):
      vars.append(var)
  return sorted(vars, key=alphanum_key)

def get_nodes(vars, port=9876):
  node = lambda k: { 'id': k, 'ip': k, 'port': port }
  return [ node(environ[v]) for v in vars ]

def generate_base_partition(nodes):
  return {
            'id': 0,
            'members': nodes
         }

def generate_extended_partitions_scaling(nodes, k, partitions=3, equal=False):
  l = deque(nodes)
  perms = []
  for i in range(1, partitions + 1):
    if equal:
      members = list(l)
    else:
      members = list(l)[:k]

    part = {
             'id': i,
             'members': members
           }
    perms.append(part)
    l.rotate(-2)
  return perms

if __name__ == '__main__':
  vars = get_OC_vars()
  # NOTE: likely prone to errors
  nodes = get_nodes(vars)
  num = None
  if len(sys.argv) >= 3:
      num = int(sys.argv[2])
      try:
          equal = "-e" in sys.argv[3]
      except:
          equal = False
  if num:
      nodes = nodes[:num]

  base_partition = generate_base_partition([v.get('id') for v in nodes])
  extended_partitions = generate_extended_partitions_scaling([v.get('id') for v in nodes],
          3, equal=equal)
  partitions = []
  partitions.append(base_partition)
  partitions.extend(extended_partitions)
  name = hash("HAScaling")
  data = {
           'name': name,
           'nodes': nodes,
           'partitions': partitions
         }
  output = json.dumps(data, indent=4)

  if len(sys.argv) >= 2:
    filename = sys.argv[1]
    with open(filename, 'w') as f:
      f.write(output)
  else:
    print output