<commit_before>#!/usr/bin/python
# Copyright 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

'''
This script provides tools to map BattOrs to phones.

Phones are identified by the following string:

"Phone serial number" - Serial number of the phone. This can be
obtained via 'adb devices' or 'usb-devices', and is not expected
to change for a given phone.

BattOrs are identified by the following two strings:

"BattOr serial number" - Serial number of the BattOr. This can be
obtained via 'usb-devices', and is not expected to change for
a given BattOr.

"BattOr path" - The path of the form '/dev/ttyUSB*' that is used
to communicate with the BattOr (the battor_agent binary takes
this BattOr path as a parameter). The BattOr path is frequently
reassigned by the OS, most often when the device is disconnected
and then reconnected. Thus, the BattOr path cannot be expected
to be stable.

In a typical application, the user will require the BattOr path
for the BattOr that is plugged into a given phone. For instance,
the user will be running tracing on a particular phone, and will
need to know which BattOr path to use to communicate with the BattOr
to get the corresponding power trace.

Getting this mapping requires two steps: (1) determining the
mapping between phone serial numbers and BattOr serial numbers, and
(2) getting the BattOr path corresponding to a given BattOr serial
number.

For step (1), we generate a JSON file giving this mapping. This
JSON file consists of a list of items of the following form:
[{'phone': <phone serial 1>, 'battor': <battor serial 1>},
{'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

The default way to generate this JSON file is using the function
GenerateSerialMapFile, which generates a mapping based on assuming
that the system has two identical USB hubs connected to it, and
the phone plugged into physical port number 1 on one hub corresponds
to the BattOr plugged into physical port number 1 on the other hub,
and similarly with physical port numbers 2, 3, etc. This generates
the map file based on the structure at the time GenerateSerialMapFile called.
Note that after the map file is generated, port numbers are no longer used;
the user could move around the devices in the ports without affecting
which phone goes with which BattOr. (Thus, if the user wanted to update the
mapping to match the new port connections, the user would have to
re-generate this file.)

The script update_mapping.py will do this updating from the command line.

If the user wanted to specify a custom mapping, the user could instead
create the JSON file manually. (In this case, hubs would not be necessary
and the physical ports connected would be irrelevant.)

Step (2) is conducted through the function GetBattorPathFromPhoneSerial,
which takes a serial number mapping generated via step (1) and a phone
serial number, then gets the corresponding BattOr serial number from the
map and determines its BattOr path (e.g. /dev/ttyUSB0). Since BattOr paths
can change if devices are connected and disconnected (even if connected
or disconnected via the same port) this function should be called to
determine the BattOr path every time before connecting to the BattOr.

Note that if there is only one BattOr connected to the system, then
GetBattorPathFromPhoneSerial will always return that BattOr and will ignore
the mapping file. Thus, if the user never has more than one BattOr connected
to the system, the user will not need to generate mapping files.
'''

import json
import collections

from battor import battor_error
from devil.utils import find_usb_devices
from devil.utils import usb_hubs

def GetBattorList(device_tree_map):
  return [x for x in find_usb_devices.GetTTYList()
          if IsBattor(x, device_tree_map)]

def IsBattor(tty_string, device_tree_map):
  (bus, device) = find_usb_devices.GetBusDeviceFromTTY(tty_string)
  node = device_tree_map[bus].FindDeviceNumber(device)
  return '0403:6001' in node.desc

def GetBattorSerialNumbers(device_tree_map):
  for x in find_usb_devices.GetTTYList():
    if IsBattor(x, device_tree_map):
      (bus, device) = find_usb_devices.GetBusDeviceFromTTY(x)
      devnode = device_tree_map[bus].FindDeviceNumber(device)
      yield devnode.serial

def ReadSerialMapFile(filename):
  """Reads JSON file giving phone-to-battor serial number map.

  Parses a JSON file consisting of a list of items of the following form:
  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
  {'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Returns dictionary of the form:

  {<phone serial 1>: <BattOr serial 1>,
   <phone serial 2>: <BattOr serial 2>}

  Args:
      filename: Name of file to read.
  """
  result = {}
  with open(filename, 'r') as infile:
    in_dict = json.load(infile)
  for x in in_dict:
    result[x['phone']] = x['battor']
  return result

def WriteSerialMapFile(filename, serial_map):
  """Writes a map of phone serial numbers to BattOr serial numbers to file.

  Writes a JSON file consisting of a list of items of the following form:
  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
  {'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Mapping is based on the physical port numbers
  of the hubs that the BattOrs and phones are connected to.

  Args:
      filename: Name of file to write.
      serial_map: Serial map {phone: battor}
  """
  result = []
  for (phone, battor) in serial_map.iteritems():
    result.append({'phone': phone, 'battor': battor})
  with open(filename, 'w') as outfile:
    json.dump(result, outfile)

def GenerateSerialMap(hub_types=None):
  """Generates a map of phone serial numbers to BattOr serial numbers.

  Generates a dict of:
  {<phone serial 1>: <battor serial 1>,
   <phone serial 2>: <battor serial 2>}
  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Mapping is based on the physical port numbers
  of the hubs that the BattOrs and phones are connected to.

  Args:
      hub_types: List of hub types to check for.
      Defaults to ['plugable_7port']
  """
  hub_types = [usb_hubs.GetHubType(x)
               for x in hub_types or ['plugable_7port']]
  devtree = find_usb_devices.GetBusNumberToDeviceTreeMap()

  # List of serial numbers in the system that represent BattOrs.
  battor_serials = list(GetBattorSerialNumbers(devtree))

  # List of dictionaries, one for each hub, that maps the physical
  # port number to the serial number of that hub. For instance, in a 2
  # hub system, this could return [{1:'ab', 2:'cd'}, {1:'jkl', 2:'xyz'}]
  # where 'ab' and 'cd' are the phone serial numbers and 'jkl' and 'xyz'
  # are the BattOr serial numbers.
  port_to_serial = find_usb_devices.GetAllPhysicalPortToSerialMaps(
      hub_types, device_tree_map=devtree)

  class serials(object):
    def __init__(self):
      self.phone = None
      self.battor = None

  # Map of {physical port number: [phone serial #, BattOr serial #]. This
  # map is populated by executing the code below. For instance, in the above
  # example, after the code below is executed, port_to_devices would equal
  # {1: ['ab', 'jkl'], 2: ['cd', 'xyz']}
  port_to_devices = collections.defaultdict(serials)
  for hub in port_to_serial:
    for (port, serial) in hub.iteritems():
      if serial in battor_serials:
        if port_to_devices[port].battor is not None:
          raise battor_error.BattorError('Multiple BattOrs on same port number')
        else:
          port_to_devices[port].battor = serial
      else:
        if port_to_devices[port].phone is not None:
          raise battor_error.BattorError('Multiple phones on same port number')
        else:
          port_to_devices[port].phone = serial

  # Turn the port_to_devices map into a map of the form
  # {phone serial number: BattOr serial number}.
  result = {}
  for pair in port_to_devices.values():
    if pair.phone is None:
      raise battor_error.BattorError(
          'BattOr detected with no corresponding phone')
    if pair.battor is None:
      raise battor_error.BattorError(
          'Phone detected with no corresponding BattOr')
    result[pair.phone] = pair.battor
  return result

def GenerateSerialMapFile(filename, hub_types=None):
  """Generates a serial map file and writes it."""
  WriteSerialMapFile(filename, GenerateSerialMap(hub_types))

def _PhoneToPathMap(serial, serial_map, devtree):
  """Maps phone serial number to TTY path, assuming serial map is provided."""
  try:
    battor_serial = serial_map[serial]
  except KeyError:
    raise battor_error.BattorError('Serial number not found in serial map.')
  for tree in devtree.values():
    for node in tree.AllNodes():
      if isinstance(node, find_usb_devices.USBDeviceNode):
        if node.serial == battor_serial:
          bus_device_to_tty = find_usb_devices.GetBusDeviceToTTYMap()
          bus_device = (node.bus_num, node.device_num)
          try:
            return bus_device_to_tty[bus_device]
          except KeyError:
            raise battor_error.BattorError(
                'Device with given serial number not a BattOr '
                '(does not have TTY path)')

def GetBattorPathFromPhoneSerial(serial, serial_map=None,
                                 serial_map_file=None):
  """Gets the TTY path (e.g. '/dev/ttyUSB0')  to communicate with the BattOr.

  (1) If serial_map is given, it is treated as a dictionary mapping
  phone serial numbers to BattOr serial numbers. This function will get the
  TTY path for the given BattOr serial number.

  (2) If serial_map_file is given, it is treated as the name of a
  phone-to-BattOr mapping file (generated with GenerateSerialMapFile)
  and this will be loaded and used as the dict to map port numbers to
  BattOr serial numbers.

  You can only give one of serial_map and serial_map_file.

  Args:
    serial: Serial number of phone connected on the same physical port that
    the BattOr is connected to.
    serial_map: Map of phone serial numbers to BattOr serial numbers, given
    as a dictionary.
    serial_map_file: Map of phone serial numbers to BattOr serial numbers,
    given as a file.
    hub_types: List of hub types to check for. Used only if serial_map_file
    is None.

  Returns:
    Device string used to communicate with device.

  Raises:
    ValueError: If serial number is not given.
    BattorError: If BattOr not found or unexpected USB topology.
  """
  # If there's only one BattOr connected to the system, just use that one.
  # This allows for use on, e.g., a developer's workstation with no hubs.
  devtree = find_usb_devices.GetBusNumberToDeviceTreeMap()
  all_battors = GetBattorList(devtree)
  if len(all_battors) == 1:
    return '/dev/' + all_battors[0]

  if not serial:
    raise battor_error.BattorError(
        'Two or more BattOrs connected, no serial provided')

  if serial_map and serial_map_file:
    raise ValueError('Cannot specify both serial_map and serial_map_file')

  if serial_map_file:
    serial_map = ReadSerialMapFile(serial_map_file)

  tty_string = _PhoneToPathMap(serial, serial_map, devtree)

  if not tty_string:
    raise battor_error.BattorError(
        'No device with given serial number detected.')

  if IsBattor(tty_string, devtree):
    return '/dev/' + tty_string
  else:
    raise battor_error.BattorError(
        'Device with given serial number is not a BattOr.')
<commit_msg>Fix single BattOr case in BattOr-device mapping script.<commit_after>#!/usr/bin/python
# Copyright 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

'''
This script provides tools to map BattOrs to phones.

Phones are identified by the following string:

"Phone serial number" - Serial number of the phone. This can be
obtained via 'adb devices' or 'usb-devices', and is not expected
to change for a given phone.

BattOrs are identified by the following two strings:

"BattOr serial number" - Serial number of the BattOr. This can be
obtained via 'usb-devices', and is not expected to change for
a given BattOr.

"BattOr path" - The path of the form '/dev/ttyUSB*' that is used
to communicate with the BattOr (the battor_agent binary takes
this BattOr path as a parameter). The BattOr path is frequently
reassigned by the OS, most often when the device is disconnected
and then reconnected. Thus, the BattOr path cannot be expected
to be stable.

In a typical application, the user will require the BattOr path
for the BattOr that is plugged into a given phone. For instance,
the user will be running tracing on a particular phone, and will
need to know which BattOr path to use to communicate with the BattOr
to get the corresponding power trace.

Getting this mapping requires two steps: (1) determining the
mapping between phone serial numbers and BattOr serial numbers, and
(2) getting the BattOr path corresponding to a given BattOr serial
number.

For step (1), we generate a JSON file giving this mapping. This
JSON file consists of a list of items of the following form:
[{'phone': <phone serial 1>, 'battor': <battor serial 1>},
{'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

The default way to generate this JSON file is using the function
GenerateSerialMapFile, which generates a mapping based on assuming
that the system has two identical USB hubs connected to it, and
the phone plugged into physical port number 1 on one hub corresponds
to the BattOr plugged into physical port number 1 on the other hub,
and similarly with physical port numbers 2, 3, etc. This generates
the map file based on the structure at the time GenerateSerialMapFile called.
Note that after the map file is generated, port numbers are no longer used;
the user could move around the devices in the ports without affecting
which phone goes with which BattOr. (Thus, if the user wanted to update the
mapping to match the new port connections, the user would have to
re-generate this file.)

The script update_mapping.py will do this updating from the command line.

If the user wanted to specify a custom mapping, the user could instead
create the JSON file manually. (In this case, hubs would not be necessary
and the physical ports connected would be irrelevant.)

Step (2) is conducted through the function GetBattorPathFromPhoneSerial,
which takes a serial number mapping generated via step (1) and a phone
serial number, then gets the corresponding BattOr serial number from the
map and determines its BattOr path (e.g. /dev/ttyUSB0). Since BattOr paths
can change if devices are connected and disconnected (even if connected
or disconnected via the same port) this function should be called to
determine the BattOr path every time before connecting to the BattOr.

Note that if there is only one BattOr connected to the system, then
GetBattorPathFromPhoneSerial will always return that BattOr and will ignore
the mapping file. Thus, if the user never has more than one BattOr connected
to the system, the user will not need to generate mapping files.
'''

import json
import collections

from battor import battor_error
from devil.utils import find_usb_devices
from devil.utils import usb_hubs

def GetBattorList(device_tree_map):
  return [x for x in find_usb_devices.GetTTYList()
          if IsBattor(x, device_tree_map)]

def IsBattor(tty_string, device_tree_map):
  (bus, device) = find_usb_devices.GetBusDeviceFromTTY(tty_string)
  node = device_tree_map[bus].FindDeviceNumber(device)
  return '0403:6001' in node.desc

def GetBattorSerialNumbers(device_tree_map):
  for x in find_usb_devices.GetTTYList():
    if IsBattor(x, device_tree_map):
      (bus, device) = find_usb_devices.GetBusDeviceFromTTY(x)
      devnode = device_tree_map[bus].FindDeviceNumber(device)
      yield devnode.serial

def ReadSerialMapFile(filename):
  """Reads JSON file giving phone-to-battor serial number map.

  Parses a JSON file consisting of a list of items of the following form:
  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
  {'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Returns dictionary of the form:

  {<phone serial 1>: <BattOr serial 1>,
   <phone serial 2>: <BattOr serial 2>}

  Args:
      filename: Name of file to read.
  """
  result = {}
  with open(filename, 'r') as infile:
    in_dict = json.load(infile)
  for x in in_dict:
    result[x['phone']] = x['battor']
  return result

def WriteSerialMapFile(filename, serial_map):
  """Writes a map of phone serial numbers to BattOr serial numbers to file.

  Writes a JSON file consisting of a list of items of the following form:
  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
  {'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Mapping is based on the physical port numbers
  of the hubs that the BattOrs and phones are connected to.

  Args:
      filename: Name of file to write.
      serial_map: Serial map {phone: battor}
  """
  result = []
  for (phone, battor) in serial_map.iteritems():
    result.append({'phone': phone, 'battor': battor})
  with open(filename, 'w') as outfile:
    json.dump(result, outfile)

def GenerateSerialMap(hub_types=None):
  """Generates a map of phone serial numbers to BattOr serial numbers.

  Generates a dict of:
  {<phone serial 1>: <battor serial 1>,
   <phone serial 2>: <battor serial 2>}
  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Mapping is based on the physical port numbers
  of the hubs that the BattOrs and phones are connected to.

  Args:
      hub_types: List of hub types to check for.
      Defaults to ['plugable_7port']
  """
  hub_types = [usb_hubs.GetHubType(x)
               for x in hub_types or ['plugable_7port']]
  devtree = find_usb_devices.GetBusNumberToDeviceTreeMap()

  # List of serial numbers in the system that represent BattOrs.
  battor_serials = list(GetBattorSerialNumbers(devtree))

  # If there's only one BattOr in the system, then a serial number ma
  # is not necessary.
  if len(battor_serials) == 1:
    return {}

  # List of dictionaries, one for each hub, that maps the physical
  # port number to the serial number of that hub. For instance, in a 2
  # hub system, this could return [{1:'ab', 2:'cd'}, {1:'jkl', 2:'xyz'}]
  # where 'ab' and 'cd' are the phone serial numbers and 'jkl' and 'xyz'
  # are the BattOr serial numbers.
  port_to_serial = find_usb_devices.GetAllPhysicalPortToSerialMaps(
      hub_types, device_tree_map=devtree)

  class serials(object):
    def __init__(self):
      self.phone = None
      self.battor = None

  # Map of {physical port number: [phone serial #, BattOr serial #]. This
  # map is populated by executing the code below. For instance, in the above
  # example, after the code below is executed, port_to_devices would equal
  # {1: ['ab', 'jkl'], 2: ['cd', 'xyz']}
  port_to_devices = collections.defaultdict(serials)
  for hub in port_to_serial:
    for (port, serial) in hub.iteritems():
      if serial in battor_serials:
        if port_to_devices[port].battor is not None:
          raise battor_error.BattorError('Multiple BattOrs on same port number')
        else:
          port_to_devices[port].battor = serial
      else:
        if port_to_devices[port].phone is not None:
          raise battor_error.BattorError('Multiple phones on same port number')
        else:
          port_to_devices[port].phone = serial

  # Turn the port_to_devices map into a map of the form
  # {phone serial number: BattOr serial number}.
  result = {}
  for pair in port_to_devices.values():
    if pair.phone is None:
      raise battor_error.BattorError(
          'BattOr detected with no corresponding phone')
    if pair.battor is None:
      raise battor_error.BattorError(
          'Phone detected with no corresponding BattOr')
    result[pair.phone] = pair.battor
  return result

def GenerateSerialMapFile(filename, hub_types=None):
  """Generates a serial map file and writes it."""
  WriteSerialMapFile(filename, GenerateSerialMap(hub_types))

def _PhoneToPathMap(serial, serial_map, devtree):
  """Maps phone serial number to TTY path, assuming serial map is provided."""
  try:
    battor_serial = serial_map[serial]
  except KeyError:
    raise battor_error.BattorError('Serial number not found in serial map.')
  for tree in devtree.values():
    for node in tree.AllNodes():
      if isinstance(node, find_usb_devices.USBDeviceNode):
        if node.serial == battor_serial:
          bus_device_to_tty = find_usb_devices.GetBusDeviceToTTYMap()
          bus_device = (node.bus_num, node.device_num)
          try:
            return bus_device_to_tty[bus_device]
          except KeyError:
            raise battor_error.BattorError(
                'Device with given serial number not a BattOr '
                '(does not have TTY path)')

def GetBattorPathFromPhoneSerial(serial, serial_map=None,
                                 serial_map_file=None):
  """Gets the TTY path (e.g. '/dev/ttyUSB0')  to communicate with the BattOr.

  (1) If serial_map is given, it is treated as a dictionary mapping
  phone serial numbers to BattOr serial numbers. This function will get the
  TTY path for the given BattOr serial number.

  (2) If serial_map_file is given, it is treated as the name of a
  phone-to-BattOr mapping file (generated with GenerateSerialMapFile)
  and this will be loaded and used as the dict to map port numbers to
  BattOr serial numbers.

  You can only give one of serial_map and serial_map_file.

  Args:
    serial: Serial number of phone connected on the same physical port that
    the BattOr is connected to.
    serial_map: Map of phone serial numbers to BattOr serial numbers, given
    as a dictionary.
    serial_map_file: Map of phone serial numbers to BattOr serial numbers,
    given as a file.
    hub_types: List of hub types to check for. Used only if serial_map_file
    is None.

  Returns:
    Device string used to communicate with device.

  Raises:
    ValueError: If serial number is not given.
    BattorError: If BattOr not found or unexpected USB topology.
  """
  # If there's only one BattOr connected to the system, just use that one.
  # This allows for use on, e.g., a developer's workstation with no hubs.
  devtree = find_usb_devices.GetBusNumberToDeviceTreeMap()
  all_battors = GetBattorList(devtree)
  if len(all_battors) == 1:
    return '/dev/' + all_battors[0]

  if not serial:
    raise battor_error.BattorError(
        'Two or more BattOrs connected, no serial provided')

  if serial_map and serial_map_file:
    raise ValueError('Cannot specify both serial_map and serial_map_file')

  if serial_map_file:
    serial_map = ReadSerialMapFile(serial_map_file)

  tty_string = _PhoneToPathMap(serial, serial_map, devtree)

  if not tty_string:
    raise battor_error.BattorError(
        'No device with given serial number detected.')

  if IsBattor(tty_string, devtree):
    return '/dev/' + tty_string
  else:
    raise battor_error.BattorError(
        'Device with given serial number is not a BattOr.')
<|endoftext|>


<commit_before>#!/usr/bin/python
# Copyright 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import json
import collections

from battor import battor_error
from devil.utils import find_usb_devices
from devil.utils import usb_hubs

def GetBattorList(device_tree_map):
  return [x for x in find_usb_devices.GetTTYList()
          if IsBattor(x, device_tree_map)]

def IsBattor(tty_string, device_tree_map):
  (bus, device) = find_usb_devices.GetBusDeviceFromTTY(tty_string)
  node = device_tree_map[bus].FindDeviceNumber(device)
  return '0403:6001' in node.desc

def GetBattorSerialNumbers(device_tree_map):
  for x in find_usb_devices.GetTTYList():
    if IsBattor(x, device_tree_map):
      (bus, device) = find_usb_devices.GetBusDeviceFromTTY(x)
      devnode = device_tree_map[bus].FindDeviceNumber(device)
      yield devnode.serial

def ReadSerialMapFile(filename):
  """Reads JSON file giving phone-to-battor serial number map.

  Parses a JSON file consisting of a list of items of the following form:
  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
  {'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Returns dictionary of the form:

  {<phone serial 1>: <BattOr serial 1>,
   <phone serial 2>: <BattOr serial 2>}

  Args:
      filename: Name of file to read.
  """
  result = {}
  with open(filename, 'r') as infile:
    in_dict = json.load(infile)
  for x in in_dict:
    result[x['phone']] = x['battor']
  return result

def WriteSerialMapFile(filename, serial_map):
  """Writes a map of phone serial numbers to BattOr serial numbers to file.

  Writes a JSON file consisting of a list of items of the following form:
  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
  {'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Mapping is based on the physical port numbers
  of the hubs that the BattOrs and phones are connected to.

  Args:
      filename: Name of file to write.
      serial_map: Serial map {phone: battor}
  """
  result = []
  for (phone, battor) in serial_map.iteritems():
    result.append({'phone': phone, 'battor': battor})
  with open(filename, 'w') as outfile:
    json.dump(result, outfile)

def GenerateSerialMap(hub_types=None):
  """Generates a map of phone serial numbers to BattOr serial numbers.

  Generates a dict of:
  {<phone serial 1>: <battor serial 1>,
   <phone serial 2>: <battor serial 2>}
  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Mapping is based on the physical port numbers
  of the hubs that the BattOrs and phones are connected to.

  Args:
      hub_types: List of hub types to check for.
      Defaults to ['plugable_7port']
  """
  hub_types = [usb_hubs.GetHubType(x)
               for x in hub_types or ['plugable_7port']]
  devtree = find_usb_devices.GetBusNumberToDeviceTreeMap()

  # List of serial numbers in the system that represent BattOrs.
  battor_serials = list(GetBattorSerialNumbers(devtree))

  # List of dictionaries, one for each hub, that maps the physical
  # port number to the serial number of that hub. For instance, in a 2
  # hub system, this could return [{1:'ab', 2:'cd'}, {1:'jkl', 2:'xyz'}]
  # where 'ab' and 'cd' are the phone serial numbers and 'jkl' and 'xyz'
  # are the BattOr serial numbers.
  port_to_serial = find_usb_devices.GetAllPhysicalPortToSerialMaps(
      hub_types, device_tree_map=devtree)

  class serials(object):
    def __init__(self):
      self.phone = None
      self.battor = None

  # Map of {physical port number: [phone serial #, BattOr serial #]. This
  # map is populated by executing the code below. For instance, in the above
  # example, after the code below is executed, port_to_devices would equal
  # {1: ['ab', 'jkl'], 2: ['cd', 'xyz']}
  port_to_devices = collections.defaultdict(serials)
  for hub in port_to_serial:
    for (port, serial) in hub.iteritems():
      if serial in battor_serials:
        if port_to_devices[port].battor is not None:
          raise battor_error.BattorError('Multiple BattOrs on same port number')
        else:
          port_to_devices[port].battor = serial
      else:
        if port_to_devices[port].phone is not None:
          raise battor_error.BattorError('Multiple phones on same port number')
        else:
          port_to_devices[port].phone = serial

  # Turn the port_to_devices map into a map of the form
  # {phone serial number: BattOr serial number}.
  result = {}
  for pair in port_to_devices.values():
    if pair.phone is None:
      raise battor_error.BattorError(
          'BattOr detected with no corresponding phone')
    if pair.battor is None:
      raise battor_error.BattorError(
          'Phone detected with no corresponding BattOr')
    result[pair.phone] = pair.battor
  return result

def GenerateSerialMapFile(filename, hub_types=None):
  """Generates a serial map file and writes it."""
  WriteSerialMapFile(filename, GenerateSerialMap(hub_types))

def _PhoneToPathMap(serial, serial_map, devtree):
  """Maps phone serial number to TTY path, assuming serial map is provided."""
  battor_serial = serial_map[serial]
  for tree in devtree.values():
    for node in tree.AllNodes():
      if isinstance(node, find_usb_devices.USBDeviceNode):
        if node.serial == battor_serial:
          bus_device_to_tty = find_usb_devices.GetBusDeviceToTTYMap()
          bus_device = (node.bus_num, node.device_num)
          try:
            return bus_device_to_tty[bus_device]
          except KeyError:
            raise battor_error.BattorError(
                'Device with given serial number not a BattOr '
                '(does not have TTY path)')

def GetBattorPathFromPhoneSerial(serial, serial_map=None,
                                 serial_map_file=None):
  """Gets the TTY path (e.g. '/dev/ttyUSB0')  to communicate with the BattOr.

  (1) If serial_map is given, it is treated as a dictionary mapping
  phone serial numbers to BattOr serial numbers. This function will get the
  TTY path for the given BattOr serial number.

  (2) If serial_map_file is given, it is treated as the name of a
  phone-to-BattOr mapping file (generated with GenerateSerialMapFile)
  and this will be loaded and used as the dict to map port numbers to
  BattOr serial numbers.

  You can only give one of serial_map and serial_map_file.

  Args:
    serial: Serial number of phone connected on the same physical port that
    the BattOr is connected to.
    serial_map: Map of phone serial numbers to BattOr serial numbers, given
    as a dictionary.
    serial_map_file: Map of phone serial numbers to BattOr serial numbers,
    given as a file.
    hub_types: List of hub types to check for. Used only if serial_map_file
    is None.

  Returns:
    Device string used to communicate with device.

  Raises:
    ValueError: If serial number is not given.
    BattorError: If BattOr not found or unexpected USB topology.
  """
  # If there's only one BattOr connected to the system, just use that one.
  # This allows for use on, e.g., a developer's workstation with no hubs.
  devtree = find_usb_devices.GetBusNumberToDeviceTreeMap()
  all_battors = GetBattorList(devtree)
  if len(all_battors) == 1:
    return '/dev/' + all_battors[0]

  if not serial:
    raise battor_error.BattorError(
        'Two or more BattOrs connected, no serial provided')

  if serial_map and serial_map_file:
    raise ValueError('Cannot specify both serial_map and serial_map_file')

  if serial_map_file:
    serial_map = ReadSerialMapFile(serial_map_file)

  tty_string = _PhoneToPathMap(serial, serial_map, devtree)

  if not tty_string:
    raise battor_error.BattorError(
        'No device with given serial number detected.')

  if IsBattor(tty_string, devtree):
    return '/dev/' + tty_string
  else:
    raise battor_error.BattorError(
        'Device with given serial number is not a BattOr.')
<commit_msg>Update documentation for BattOr device mapping.<commit_after>#!/usr/bin/python
# Copyright 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

'''
This script provides tools to map BattOrs to phones.

Phones are identified by the following string:

"Phone serial number" - Serial number of the phone. This can be
obtained via 'adb devices' or 'usb-devices', and is not expected
to change for a given phone.

BattOrs are identified by the following two strings:

"BattOr serial number" - Serial number of the BattOr. This can be
obtained via 'usb-devices', and is not expected to change for
a given BattOr.

"BattOr path" - The path of the form '/dev/ttyUSB*' that is used
to communicate with the BattOr (the battor_agent binary takes
this BattOr path as a parameter). The BattOr path is frequently
reassigned by the OS, most often when the device is disconnected
and then reconnected. Thus, the BattOr path cannot be expected
to be stable.

In a typical application, the user will require the BattOr path
for the BattOr that is plugged into a given phone. For instance,
the user will be running tracing on a particular phone, and will
need to know which BattOr path to use to communicate with the BattOr
to get the corresponding power trace.

Getting this mapping requires two steps: (1) determining the
mapping between phone serial numbers and BattOr serial numbers, and
(2) getting the BattOr path corresponding to a given BattOr serial
number.

For step (1), we generate a JSON file giving this mapping. This
JSON file consists of a list of items of the following form:
[{'phone': <phone serial 1>, 'battor': <battor serial 1>},
{'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

The default way to generate this JSON file is using the function
GenerateSerialMapFile, which generates a mapping based on assuming
that the system has two identical USB hubs connected to it, and
the phone plugged into physical port number 1 on one hub corresponds
to the BattOr plugged into physical port number 1 on the other hub,
and similarly with physical port numbers 2, 3, etc. This generates
the map file based on the structure at the time GenerateSerialMapFile called.
Note that after the map file is generated, port numbers are no longer used;
the user could move around the devices in the ports without affecting
which phone goes with which BattOr. (Thus, if the user wanted to update the
mapping to match the new port connections, the user would have to
re-generate this file.)

The script update_mapping.py will do this updating from the command line.

If the user wanted to specify a custom mapping, the user could instead
create the JSON file manually. (In this case, hubs would not be necessary
and the physical ports connected would be irrelevant.)

Step (2) is conducted through the function GetBattorPathFromPhoneSerial,
which takes a serial number mapping generated via step (1) and a phone
serial number, then gets the corresponding BattOr serial number from the
map and determines its BattOr path (e.g. /dev/ttyUSB0). Since BattOr paths
can change if devices are connected and disconnected (even if connected
or disconnected via the same port) this function should be called to
determine the BattOr path every time before connecting to the BattOr.

Note that if there is only one BattOr connected to the system, then
GetBattorPathFromPhoneSerial will always return that BattOr and will ignore
the mapping file. Thus, if the user never has more than one BattOr connected
to the system, the user will not need to generate mapping files.
'''

import json
import collections

from battor import battor_error
from devil.utils import find_usb_devices
from devil.utils import usb_hubs

def GetBattorList(device_tree_map):
  return [x for x in find_usb_devices.GetTTYList()
          if IsBattor(x, device_tree_map)]

def IsBattor(tty_string, device_tree_map):
  (bus, device) = find_usb_devices.GetBusDeviceFromTTY(tty_string)
  node = device_tree_map[bus].FindDeviceNumber(device)
  return '0403:6001' in node.desc

def GetBattorSerialNumbers(device_tree_map):
  for x in find_usb_devices.GetTTYList():
    if IsBattor(x, device_tree_map):
      (bus, device) = find_usb_devices.GetBusDeviceFromTTY(x)
      devnode = device_tree_map[bus].FindDeviceNumber(device)
      yield devnode.serial

def ReadSerialMapFile(filename):
  """Reads JSON file giving phone-to-battor serial number map.

  Parses a JSON file consisting of a list of items of the following form:
  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
  {'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Returns dictionary of the form:

  {<phone serial 1>: <BattOr serial 1>,
   <phone serial 2>: <BattOr serial 2>}

  Args:
      filename: Name of file to read.
  """
  result = {}
  with open(filename, 'r') as infile:
    in_dict = json.load(infile)
  for x in in_dict:
    result[x['phone']] = x['battor']
  return result

def WriteSerialMapFile(filename, serial_map):
  """Writes a map of phone serial numbers to BattOr serial numbers to file.

  Writes a JSON file consisting of a list of items of the following form:
  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
  {'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]

  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Mapping is based on the physical port numbers
  of the hubs that the BattOrs and phones are connected to.

  Args:
      filename: Name of file to write.
      serial_map: Serial map {phone: battor}
  """
  result = []
  for (phone, battor) in serial_map.iteritems():
    result.append({'phone': phone, 'battor': battor})
  with open(filename, 'w') as outfile:
    json.dump(result, outfile)

def GenerateSerialMap(hub_types=None):
  """Generates a map of phone serial numbers to BattOr serial numbers.

  Generates a dict of:
  {<phone serial 1>: <battor serial 1>,
   <phone serial 2>: <battor serial 2>}
  indicating which phone serial numbers should be matched with
  which BattOr serial numbers. Mapping is based on the physical port numbers
  of the hubs that the BattOrs and phones are connected to.

  Args:
      hub_types: List of hub types to check for.
      Defaults to ['plugable_7port']
  """
  hub_types = [usb_hubs.GetHubType(x)
               for x in hub_types or ['plugable_7port']]
  devtree = find_usb_devices.GetBusNumberToDeviceTreeMap()

  # List of serial numbers in the system that represent BattOrs.
  battor_serials = list(GetBattorSerialNumbers(devtree))

  # List of dictionaries, one for each hub, that maps the physical
  # port number to the serial number of that hub. For instance, in a 2
  # hub system, this could return [{1:'ab', 2:'cd'}, {1:'jkl', 2:'xyz'}]
  # where 'ab' and 'cd' are the phone serial numbers and 'jkl' and 'xyz'
  # are the BattOr serial numbers.
  port_to_serial = find_usb_devices.GetAllPhysicalPortToSerialMaps(
      hub_types, device_tree_map=devtree)

  class serials(object):
    def __init__(self):
      self.phone = None
      self.battor = None

  # Map of {physical port number: [phone serial #, BattOr serial #]. This
  # map is populated by executing the code below. For instance, in the above
  # example, after the code below is executed, port_to_devices would equal
  # {1: ['ab', 'jkl'], 2: ['cd', 'xyz']}
  port_to_devices = collections.defaultdict(serials)
  for hub in port_to_serial:
    for (port, serial) in hub.iteritems():
      if serial in battor_serials:
        if port_to_devices[port].battor is not None:
          raise battor_error.BattorError('Multiple BattOrs on same port number')
        else:
          port_to_devices[port].battor = serial
      else:
        if port_to_devices[port].phone is not None:
          raise battor_error.BattorError('Multiple phones on same port number')
        else:
          port_to_devices[port].phone = serial

  # Turn the port_to_devices map into a map of the form
  # {phone serial number: BattOr serial number}.
  result = {}
  for pair in port_to_devices.values():
    if pair.phone is None:
      raise battor_error.BattorError(
          'BattOr detected with no corresponding phone')
    if pair.battor is None:
      raise battor_error.BattorError(
          'Phone detected with no corresponding BattOr')
    result[pair.phone] = pair.battor
  return result

def GenerateSerialMapFile(filename, hub_types=None):
  """Generates a serial map file and writes it."""
  WriteSerialMapFile(filename, GenerateSerialMap(hub_types))

def _PhoneToPathMap(serial, serial_map, devtree):
  """Maps phone serial number to TTY path, assuming serial map is provided."""
  battor_serial = serial_map[serial]
  for tree in devtree.values():
    for node in tree.AllNodes():
      if isinstance(node, find_usb_devices.USBDeviceNode):
        if node.serial == battor_serial:
          bus_device_to_tty = find_usb_devices.GetBusDeviceToTTYMap()
          bus_device = (node.bus_num, node.device_num)
          try:
            return bus_device_to_tty[bus_device]
          except KeyError:
            raise battor_error.BattorError(
                'Device with given serial number not a BattOr '
                '(does not have TTY path)')

def GetBattorPathFromPhoneSerial(serial, serial_map=None,
                                 serial_map_file=None):
  """Gets the TTY path (e.g. '/dev/ttyUSB0')  to communicate with the BattOr.

  (1) If serial_map is given, it is treated as a dictionary mapping
  phone serial numbers to BattOr serial numbers. This function will get the
  TTY path for the given BattOr serial number.

  (2) If serial_map_file is given, it is treated as the name of a
  phone-to-BattOr mapping file (generated with GenerateSerialMapFile)
  and this will be loaded and used as the dict to map port numbers to
  BattOr serial numbers.

  You can only give one of serial_map and serial_map_file.

  Args:
    serial: Serial number of phone connected on the same physical port that
    the BattOr is connected to.
    serial_map: Map of phone serial numbers to BattOr serial numbers, given
    as a dictionary.
    serial_map_file: Map of phone serial numbers to BattOr serial numbers,
    given as a file.
    hub_types: List of hub types to check for. Used only if serial_map_file
    is None.

  Returns:
    Device string used to communicate with device.

  Raises:
    ValueError: If serial number is not given.
    BattorError: If BattOr not found or unexpected USB topology.
  """
  # If there's only one BattOr connected to the system, just use that one.
  # This allows for use on, e.g., a developer's workstation with no hubs.
  devtree = find_usb_devices.GetBusNumberToDeviceTreeMap()
  all_battors = GetBattorList(devtree)
  if len(all_battors) == 1:
    return '/dev/' + all_battors[0]

  if not serial:
    raise battor_error.BattorError(
        'Two or more BattOrs connected, no serial provided')

  if serial_map and serial_map_file:
    raise ValueError('Cannot specify both serial_map and serial_map_file')

  if serial_map_file:
    serial_map = ReadSerialMapFile(serial_map_file)

  tty_string = _PhoneToPathMap(serial, serial_map, devtree)

  if not tty_string:
    raise battor_error.BattorError(
        'No device with given serial number detected.')

  if IsBattor(tty_string, devtree):
    return '/dev/' + tty_string
  else:
    raise battor_error.BattorError(
        'Device with given serial number is not a BattOr.')
<|endoftext|>


<commit_before>from ApplianceTrace import ApplianceTrace

import sqlalchemy
import pandas as pd

class PecanStreetDatasetAdapter():

    schema_names = {'curated': '\"PecanStreet_CuratedSets\"',
                   'raw':     '\"PecanStreet_RawData\"',
                   'shared'  :'\"PecanStreet_SharedData\"'}

    time_columns = {'curated': 'utc_15min',
                    'raw':     'localminute15minute',
                    'shared':  'localminute'}

    invalid_columns = {'curated': ['id', 'utc_15min'],
                             'raw':     ['localminute15minute'],
                             'shared':  ['localminute']}

    def __init__(self,db_url):
        '''
        Initialize an adapter using a database url_string.
        Consider the following example:
        db_url="postgresql"+"://"+user_name+":"+ps+"@"+host+":"+port+"/"+db
        '''
        self.eng = sqlalchemy.create_engine(db_url)
        self.source = "PecanStreet"

    def set_table_names(self,schema): #TODO Change this func name to "get_table_names"
        '''
        Returns a list of tables in the schema.
        '''
        df = self.get_dataframe('select * from information_schema.tables')
        df = df.groupby(['table_schema','table_name'])
        groups = [group for group in df.groups]
        table_names = [t for (s,t) in groups if s == schema_names[schema]]
        return table_names

    def verify_same_range(self):
        '''
        Check that all data points have the same range
        '''
        pass

    def get_meta_table(self,schema,table): # TODO change this func name to "get_table_metadata"
        '''
        Returns a tuple where the first element is a list of data ids for this
        schema.table and the second element is a list of the appliances
        included in this schema.table
        '''
        q = 'select distinct dataid from {}.{}'.format(schema_names[schema],table)
        result = self.eng.execute(q)
        ids = result.fetchall()
        q = 'select * from {}.{} where dataid={}'.format(schema_names[schema],table,ids[0][0])
        result = self.eng.execute(q)
        apps = result.keys()
        ids= [a[0] for a in ids]
        apps = [str(a) for a in apps ]
        return [ids,apps]

    def get_unique_dataids(self,schema,month,year,group=None):
        '''
        Returns a list of dataids for a specifc schema ("curated","shared", or
        "raw"), month (int), year (int), and group (int).
        '''
        if schema == "curated":
            schema_name = schema_names[schema]
            query = 'select distinct dataid from {0}.group{1}_disaggregated_{2}_{3:02d}'.format(schema_name,group,year,month)
            df = self.get_dataframe(query)
            return list(df["dataid"])
        elif schema == "shared":
            raise NotImplementedError
        elif schema == "raw":
            raise NotImplementedError
        else:
            raise SchemaError(schema)

    def get_month_traces(self,schema,year,month,dataid,group=None,sampling_rate="15T"):
        '''
        Returns a month-long traces for the specified month and sampling rate. Specify
        sampling rate using pd offset aliases (Ex. 15 mins -> "15T")
        '''
        if schema == "curated": # Lowest possible sampling rate is 15T
            # load dataframe and fill with zeros
            schema_name = schema_names[schema]
            query = 'select * from {0}.group{1}_disaggregated_{2}_{3:02d} where dataid={4}'.format(schema_name,group,year,month,dataid)
            df = self.get_dataframe(query).fillna(0)

            # column name for a trace series DatetimeIndex should be "time"
            df.rename(columns={'utc_15min': 'time'}, inplace=True)
            df.index = df['time'].apply(pd.to_datetime)

            # drop unneded columns
            df = df.drop(['id','dataid','time'], axis=1)

            # resample if necessary
            if not (sampling_rate == '15T' or sampling_rate == '15Min'):
                how = {col:'sum' for col in dataframe.columns}
                df = df.resample(sampling_rate, how=how)
        elif schema == "shared":
            pass
        elif schema == "raw":
            raise NotImplementedError
        else:
            raise SchemaError(schema)

        # make traces for each column
        traces = []
        for column, series in df.iteritems():
            traces.append(ApplianceTrace(series,self.source))
        return traces

    def time_align():
        '''Checks that for all traces in a home the total time lengths are the same'''
        pass

    def clean_dataframe(self,df,schema,drop_cols): # TODO update this to use "curated" "shared" or "raw" instead of full frame name
        '''
        Cleans a dataframe queried directly from the database.
        '''
        # change the time column name
        df = df.rename(columns={time_colums[schema]: 'time'})

        # use a DatetimeIndex
        df['time'] = pd.to_datetime(df['time'], format='%d/%m/%Y %H:%M:%S')
        df.set_index('time', inplace=True)

        # get some info about times
        start_time = df['time'][0]
        end_time = df['time'][-1]
        step_size = df['time'][1]-start_time # will error out if we only have one time point
        times = (start_time, end_time, step_size)

        # drop unnecessary columns
        df = df.drop(['dataid'], axis=1)
        if schema == 'curated':
            df = df.drop(['id'], axis=1)
        if len(drop_cols)!=0:
            df= df.drop(drop_cols,axis=1)

        return df, times

    def check_sample_rate(self,schema,sampling_rate):
        ##get from the data directly not like this
        accepted_rates = {'curated':'15T' ,'raw':'15' ,'shared':'1T' }

    def get_month_traces_per_dataid(self,schema,table,dataid):
        # TODO change this name
        if schema not in ['curated','raw','shared']:
            raise SchemaError(schema)
        schema_name = schema_names[schema]
        query = 'select * from {0}.{1} where dataid={2}'.format(schema_name, table, dataid)
        # TODO NEED TO CHANGE IDS
        # TODO error checking that query worked
        df = self.get_dataframe(query).fillna(0)

        df,times = self.clean_dataframe(df, schema,[])
        traces = []
        for col in df.columns:
            if not col in invalid_columns[schema]:
                meta={'source':self.source,'schema':schema,'table':table ,'dataid':dataid, 'start_time': times[0],'end_time':times[1], 'step_size':times[2] }
                traces.append(ApplianceTrace(df[col],meta))
        return traces

    def get_single_app_trace_need_house_id(self,house_df, app):
        '''by house is fastest also have get all apps below'''
        pass

    def get_app_traces_all(self,schema,table,app):
        schema_name = schema_names[schema]
        query= 'select {2} from {0}.{1}'.format(schema_name,table,app)
        df=self.get_dataframe(query)
        # TODO - does this need to be cleaned differently?

    def get_dataframe(self,query):
        '''
        Returns a Pandas dataframe with the query results
        '''
        eng_object = self.eng.execute(query)
        df = pd.DataFrame.from_records(eng_object.fetchall())
        df.columns = eng_object.keys()
        return df

    def get_list(self,query):
        '''
        Returns a list with query results.
        '''
        result  = self.eng.execute(query).fetchall()
        return [result,result.keys()]

class SchemaError(Exception):
    '''Exception raised for errors in the schema.

        Attributes:
            schema  -- nonexistent schema
    '''
    def __init__(self,schema):
        self.schema = schema

    def __str__(self):
        return "Schema {} not supported or nonexistent.".format(self.schema)
<commit_msg>update set_table_names to get_table_names<commit_after>from ApplianceTrace import ApplianceTrace

import sqlalchemy
import pandas as pd

class PecanStreetDatasetAdapter():

    schema_names =    {'curated': '\"PecanStreet_CuratedSets\"',
                       'raw':     '\"PecanStreet_RawData\"',
                       'shared':  '\"PecanStreet_SharedData\"'}

    time_columns =    {'curated': 'utc_15min',
                       'raw':     'localminute15minute',
                       'shared':  'localminute'}

    invalid_columns = {'curated': ['id', 'utc_15min'],
                       'raw':     ['localminute15minute'],
                       'shared':  ['localminute']}

    def __init__(self,db_url):
        '''
        Initialize an adapter using a database url_string.
        Consider the following example:
        db_url="postgresql://user_name:password@host.url:port/db"
        '''
        self.eng = sqlalchemy.create_engine(db_url)
        self.source = "PecanStreet"

    def get_table_names(self,schema):
        '''
        Returns a list of tables in the schema.
        '''
        df = self.get_dataframe('select * from information_schema.tables')
        df = df.groupby(['table_schema','table_name'])
        groups = [group for group in df.groups]
        table_names = [t for (s,t) in groups if s == schema_names[schema]]
        return table_names

    def verify_same_range(self):
        '''
        Check that all data points have the same range
        '''
        pass

    def get_meta_table(self,schema,table): # TODO change this func name to "get_table_metadata"
        '''
        Returns a tuple where the first element is a list of data ids for this
        schema.table and the second element is a list of the appliances
        included in this schema.table
        '''
        q = 'select distinct dataid from {}.{}'.format(schema_names[schema],table)
        result = self.eng.execute(q)
        ids = result.fetchall()
        q = 'select * from {}.{} where dataid={}'.format(schema_names[schema],table,ids[0][0])
        result = self.eng.execute(q)
        apps = result.keys()
        ids= [a[0] for a in ids]
        apps = [str(a) for a in apps ]
        return [ids,apps]

    def get_unique_dataids(self,schema,month,year,group=None):
        '''
        Returns a list of dataids for a specifc schema ("curated","shared", or
        "raw"), month (int), year (int), and group (int).
        '''
        if schema == "curated":
            schema_name = schema_names[schema]
            query = 'select distinct dataid from {0}.group{1}_disaggregated_{2}_{3:02d}'.format(schema_name,group,year,month)
            df = self.get_dataframe(query)
            return list(df["dataid"])
        elif schema == "shared":
            raise NotImplementedError
        elif schema == "raw":
            raise NotImplementedError
        else:
            raise SchemaError(schema)

    def get_month_traces(self,schema,year,month,dataid,group=None,sampling_rate="15T"):
        '''
        Returns a month-long traces for the specified month and sampling rate. Specify
        sampling rate using pd offset aliases (Ex. 15 mins -> "15T")
        '''
        if schema == "curated": # Lowest possible sampling rate is 15T
            # load dataframe and fill with zeros
            schema_name = schema_names[schema]
            query = 'select * from {0}.group{1}_disaggregated_{2}_{3:02d} where dataid={4}'.format(schema_name,group,year,month,dataid)
            df = self.get_dataframe(query).fillna(0)

            # column name for a trace series DatetimeIndex should be "time"
            df.rename(columns={'utc_15min': 'time'}, inplace=True)
            df.index = df['time'].apply(pd.to_datetime)

            # drop unneded columns
            df = df.drop(['id','dataid','time'], axis=1)

            # resample if necessary
            if not (sampling_rate == '15T' or sampling_rate == '15Min'):
                how = {col:'sum' for col in dataframe.columns}
                df = df.resample(sampling_rate, how=how)
        elif schema == "shared":
            pass
        elif schema == "raw":
            raise NotImplementedError
        else:
            raise SchemaError(schema)

        # make traces for each column
        traces = []
        for column, series in df.iteritems():
            traces.append(ApplianceTrace(series,self.source))
        return traces

    def time_align():
        '''Checks that for all traces in a home the total time lengths are the same'''
        pass

    def clean_dataframe(self,df,schema,drop_cols): # TODO update this to use "curated" "shared" or "raw" instead of full frame name
        '''
        Cleans a dataframe queried directly from the database.
        '''
        # change the time column name
        df = df.rename(columns={time_colums[schema]: 'time'})

        # use a DatetimeIndex
        df['time'] = pd.to_datetime(df['time'], format='%d/%m/%Y %H:%M:%S')
        df.set_index('time', inplace=True)

        # get some info about times
        start_time = df['time'][0]
        end_time = df['time'][-1]
        step_size = df['time'][1]-start_time # will error out if we only have one time point
        times = (start_time, end_time, step_size)

        # drop unnecessary columns
        df = df.drop(['dataid'], axis=1)
        if schema == 'curated':
            df = df.drop(['id'], axis=1)
        if len(drop_cols)!=0:
            df= df.drop(drop_cols,axis=1)

        return df, times

    def check_sample_rate(self,schema,sampling_rate):
        ##get from the data directly not like this
        accepted_rates = {'curated':'15T' ,'raw':'15' ,'shared':'1T' }

    def get_month_traces_per_dataid(self,schema,table,dataid):
        # TODO change this name
        if schema not in ['curated','raw','shared']:
            raise SchemaError(schema)
        schema_name = schema_names[schema]
        query = 'select * from {0}.{1} where dataid={2}'.format(schema_name, table, dataid)
        # TODO NEED TO CHANGE IDS
        # TODO error checking that query worked
        df = self.get_dataframe(query).fillna(0)

        df,times = self.clean_dataframe(df, schema,[])
        traces = []
        for col in df.columns:
            if not col in invalid_columns[schema]:
                meta={'source':self.source,'schema':schema,'table':table ,'dataid':dataid, 'start_time': times[0],'end_time':times[1], 'step_size':times[2] }
                traces.append(ApplianceTrace(df[col],meta))
        return traces

    def get_single_app_trace_need_house_id(self,house_df, app):
        '''by house is fastest also have get all apps below'''
        pass

    def get_app_traces_all(self,schema,table,app):
        schema_name = schema_names[schema]
        query= 'select {2} from {0}.{1}'.format(schema_name,table,app)
        df=self.get_dataframe(query)
        # TODO - does this need to be cleaned differently?

    def get_dataframe(self,query):
        '''
        Returns a Pandas dataframe with the query results
        '''
        eng_object = self.eng.execute(query)
        df = pd.DataFrame.from_records(eng_object.fetchall())
        df.columns = eng_object.keys()
        return df

class SchemaError(Exception):
    '''Exception raised for errors in the schema.

        Attributes:
            schema  -- nonexistent schema
    '''
    def __init__(self,schema):
        self.schema = schema

    def __str__(self):
        return "Schema {} not supported or nonexistent.".format(self.schema)
<|endoftext|>


<commit_before>"""
Parse disco_aws alarm configuration.
"""

import logging
import copy

from boto.ec2.cloudwatch import MetricAlarm

from . import DiscoAutoscale
from .disco_config import read_config
from .exceptions import AlarmConfigError
from .disco_aws_util import is_truthy
from .disco_elb import DiscoELB

logger = logging.getLogger(__name__)

NOTIFICATION_LEVELS = ["critical", "info"]
NOTIFICATION_SECTION_NAME = "notifications"
DEFAULT_SECTION_NAME = 'defaults'
DEFAULT_CONFIG_FILE = "disco_alarms.ini"

class DiscoNotification(object):
    """
    Representation of a notification type, which is implemented by an SNS topic on the AWS side
    """
    def __init__(self, name, endpoints):
        self.name = name
        self.endpoints = endpoints

    def __repr__(self):
        return "Notification:{0}:{1}".format(self.name, self.endpoints)

class DiscoAlarmConfig(object):
    """
    Representation of a single alarm, and all its configuration options.
    """

    # We take in a lot of options for configuring the alarms.
    # pylint: disable=too-many-instance-attributes
    def __init__(self, options, maximum=True):
        self.namespace = options["namespace"]
        self.metric_name = options["metric_name"]
        self.hostclass = options["hostclass"]
        self.environment = options["environment"]
        self.team = options["team"]

        self.duration = int(options["duration"])
        self.period = int(options["period"])
        self.statistic = options["statistic"]
        self.custom_metric = options["custom_metric"].lower() == "true"
        self.log_pattern_metric = is_truthy(options.get("log_pattern_metric", ""))

        if "autoscaling_group_name" in options:
            self.autoscaling_group_name = options["autoscaling_group_name"]
        else:
            self.autoscaling_group_name = None

        if "es_domain_name" in options:
            self.es_domain_name = options["es_domain_name"]
        else:
            self.es_domain_name = None

        if "es_client_id" in options:
            self.es_client_id = options["es_client_id"]
        else:
            self.es_client_id = None

        if options["level"] not in NOTIFICATION_LEVELS:
            raise AlarmConfigError(
                "Level for {0} must be one of {1}"
                .format(self.metric_name, NOTIFICATION_LEVELS)
            )
        self._level = options["level"]

        if maximum:
            logger.debug("threshold_max: %s", options["threshold_max"])
            self.threshold = int(options["threshold_max"])
            self.threshold_type = "max"
        else:
            logger.debug("threshold_min: %s", options["threshold_min"])
            self.threshold = int(options["threshold_min"])
            self.threshold_type = "min"

    @property
    def dimensions(self):
        """
        Alarm Metric Dimensions
        """

        # metrics created from logs do not have dimensions
        if self.log_pattern_metric:
            return {}

        # RDS dimensions
        # hostclass name is the DBInstanceIdentifier
        if self.namespace == 'AWS/RDS':
            key = "DBInstanceIdentifier"
            # Format is {vpc}-{DbInstanceId} Underscore is not allowed in RDS instance name
            value = "-".join([self.environment, self.hostclass])
            return {key: value}

        if self.namespace == 'AWS/ELB':
            return {'LoadBalancerName': DiscoELB.get_elb_id(self.environment, self.hostclass)}

        if self.namespace == 'AWS/ES':
            return {
                'DomainName': self.es_domain_name,
                'ClientId': self.es_client_id
            }

        if self.custom_metric:
            key = "env_hostclass"
            value = "_".join([self.environment, self.hostclass])
        else:
            key = "AutoScalingGroupName"
            value = self.autoscaling_group_name
        return {key: value}

    @property
    def notification_topic(self):
        """
        The name of the topic that the alarm should route to.
        """
        return "_".join([self.team, self.environment, self._level])

    @property
    def name(self):
        """
        Alarm name
        """
        return "_".join([self.team, self.environment, self.hostclass, self.metric_name, self.threshold_type])

    @staticmethod
    def decode_alarm_name(name):
        """
        Decodes an alarm name into a dictionary with team, env, hostclass, metric_name and
        threshold_type.

        Raises AlarmConfigError if unable to parse the alarm name.
        """
        parts = name.split("_")
        if len(parts) >= 5:
            # alarm names that have extra underscores
            # Assume this means the metric name portion has underscores
            # FIXME this won't work if the alarm is missing a team name AND the metric name has underscores
            # FIXME consider using different delimiters in the future
            return {
                "team": parts[0],
                "env": parts[1],
                "hostclass": parts[2],
                # some metric names have '_'. Assume that any extra '_' in the alarm name are the metric name
                "metric_name": '_'.join(parts[3:-1]),
                "threshold_type": parts[-1],
            }
        elif len(parts) == 4:
            # alarm names that don't contain a team name and a metric name without underscores
            return {
                "team": None,
                "env": parts[0],
                "hostclass": parts[1],
                "metric_name": parts[2],
                "threshold_type": parts[3],
            }
        raise AlarmConfigError("Failed to decode alarm name {0}".format(name))

    def __repr__(self):
        return "Alarm:{0}.{1}".format(self.name, self.threshold)

    def to_metric_alarm(self, policy_arn):
        """
        Returns a MetricAlarm for a given policy.
        """
        return MetricAlarm(
            alarm_actions=[policy_arn],
            ok_actions=[policy_arn],
            comparison='>' if self.threshold_type == "max" else "<",
            dimensions=self.dimensions,
            evaluation_periods=self.duration,
            metric=self.metric_name,
            name=self.name,
            namespace=self.namespace,
            period=self.period,
            statistic=self.statistic,
            threshold=self.threshold,
        )

class DiscoAlarmsConfig(object):
    """
    Represenation of all alarms in config file.
    """

    def __init__(self, environment, config_file=None, autoscale=None, elasticsearch=None):
        self.config = read_config(config_file=(config_file or DEFAULT_CONFIG_FILE), environment=environment)
        self.environment = environment
        self.elasticsearch = elasticsearch or None
        self._autoscale = autoscale or None  # laziliy intialized
        self._defaults = None

    @property
    def defaults(self):
        """
        get default options for whole disco_alarm_config
        """
        self._defaults = self._defaults or self.get_defaults()
        return self._defaults

    @property
    def autoscale(self):
        """
        returns a lazily created disco autoscaling object
        """
        if not self._autoscale:
            self._autoscale = DiscoAutoscale(environment_name=self.environment)
        return self._autoscale

    @staticmethod
    def _decode_section_name(section):
        """
        Extract special alarm parameters which are stored in section name
        """
        section_segments = section.split(".")
        # If a section has four or more segments, treat the middle segments as the metric name and
        # the ending segment as the hostclass. This is to support some Elasticsearch metrics like
        # 'ClusterStatus.red'
        if len(section_segments) >= 4:
            team = section_segments[0]
            namespace = section_segments[1]
            metric_name = '.'.join(section_segments[2:-1])
            section_hostclass = section_segments[-1]
        elif len(section_segments) == 3:
            team, namespace, metric_name = section_segments
            section_hostclass = None
        else:
            raise AlarmConfigError("Skipping non-alarm like config section {0}.".format(section))
        return team, namespace, metric_name, section_hostclass

    # Our alarms get a little complicated, so theres more than a few variables involved...
    # pylint: disable=too-many-locals
    def _get_alarm_specification_dict(self, team, namespace, metric_name, hostclass, hostclass_specific,
                                      autoscaling_group_name=None):
        """
        Return options for alarm, inheriting values from parent.
        """
        options = copy.copy(self.defaults)
        parent_section_name = ".".join([team, namespace, metric_name])
        child_section_name = ".".join([team, namespace, metric_name, hostclass])

        if hostclass_specific:
            if self.config.has_section(parent_section_name):
                parent_options = dict(self.config.items(parent_section_name))
                options.update(parent_options)
            child_options = dict(self.config.items(child_section_name))
            options.update(child_options)
        else:
            if self.config.has_section(child_section_name):
                # The alarm will be created with the more specific hostclass config
                return None
            parent_options = dict(self.config.items(parent_section_name))
            options.update(parent_options)

        options["team"] = team
        options["namespace"] = namespace
        options["metric_name"] = metric_name
        options["hostclass"] = hostclass
        options["environment"] = self.environment

        # If an autoscaling group name was provided, include it
        if autoscaling_group_name:
            options["autoscaling_group_name"] = autoscaling_group_name
        # Otherwise, if using the EC2 namespace, we should lookup the autoscaling group name of the hostclass
        # and use that.
        elif namespace == "AWS/EC2":
            group = self.autoscale.get_existing_group(hostclass=hostclass)
            options["autoscaling_group_name"] = group.name

        if namespace == "AWS/ES" and self.elasticsearch:
            domain_name = self.elasticsearch.get_domain_name(hostclass)
            options["es_domain_name"] = domain_name
            options["es_client_id"] = self.elasticsearch.get_client_id(domain_name)

        # metrics for log files have special environment specific namespaces and hostclass specific names
        if is_truthy(options.get("log_pattern_metric")):
            options["namespace"] = namespace + '/' + self.environment
            options['metric_name'] = hostclass + '-' + metric_name

        return options

    def get_alarms(self, hostclass, autoscaling_group_name=None):
        """
        Returns list of DiscoAlarm objects for all the alarms associated with a hostclass.
        """
        alarms = []
        for section in self.config.sections():
            try:
                team, namespace, metric_name, section_hostclass = DiscoAlarmsConfig._decode_section_name(
                    section
                )
            except AlarmConfigError:
                if section in [NOTIFICATION_SECTION_NAME, DEFAULT_SECTION_NAME]:
                    continue
                else:
                    raise
            if section_hostclass and section_hostclass != hostclass:
                # Hostclass specific section for different hostclass
                continue

            if namespace == "AWS/ES" and not self.elasticsearch:
                logger.info("Skipping creation of %s.%s.%s.%s because no Elasticsearch object was provided",
                            team, namespace, metric_name, section_hostclass)
                continue

            options = self._get_alarm_specification_dict(
                team, namespace, metric_name, hostclass, section_hostclass == hostclass,
                autoscaling_group_name=autoscaling_group_name
            )

            if not options:
                # No alarms options means we don't need to make the alarm
                continue

            for threshold in ["threshold_max", "threshold_min"]:
                if threshold in options:
                    if options[threshold].isdigit():
                        alarms.append(
                            DiscoAlarmConfig(options, maximum=threshold == "threshold_max")
                        )
                    else:
                        raise AlarmConfigError(
                            "Not a valid threshold value for {0}: {1}".format(threshold, options))
        return alarms

    def get_notifications(self):
        """
        Returns list of DiscoNotification objects for all notifications in the current environment
        for the current pagerduty group.
        """
        notifications = []
        options = dict(self.config.items(NOTIFICATION_SECTION_NAME))

        for notification_name, endpoint_str in options.iteritems():
            notification_parts = notification_name.split("_")
            if len(notification_parts) < 3:
                raise AlarmConfigError("Underscores must separate team name, "
                                       "env name and notification level: %s" % notification_name)
            level = notification_parts[-1]
            env = notification_parts[1]

            if env != self.environment:
                continue
            if level not in NOTIFICATION_LEVELS:
                raise AlarmConfigError("Unsupported notification level: %s" % level)

            endpoints = endpoint_str.split(",")
            notifications.append(DiscoNotification(notification_name, endpoints))

        return notifications

    def get_defaults(self):
        """
        return dictionary of default values that will be used
        """
        return dict(self.config.items(DEFAULT_SECTION_NAME)) or {}
<commit_msg>work around the asiaq crash on unexpected alarm configs<commit_after>"""
Parse disco_aws alarm configuration.
"""

import logging
import copy

from boto.ec2.cloudwatch import MetricAlarm

from . import DiscoAutoscale
from .disco_config import read_config
from .exceptions import AlarmConfigError
from .disco_aws_util import is_truthy
from .disco_elb import DiscoELB

logger = logging.getLogger(__name__)

NOTIFICATION_LEVELS = ["critical", "info"]
NOTIFICATION_SECTION_NAME = "notifications"
DEFAULT_SECTION_NAME = 'defaults'
DEFAULT_CONFIG_FILE = "disco_alarms.ini"

class DiscoNotification(object):
    """
    Representation of a notification type, which is implemented by an SNS topic on the AWS side
    """
    def __init__(self, name, endpoints):
        self.name = name
        self.endpoints = endpoints

    def __repr__(self):
        return "Notification:{0}:{1}".format(self.name, self.endpoints)

class DiscoAlarmConfig(object):
    """
    Representation of a single alarm, and all its configuration options.
    """

    # We take in a lot of options for configuring the alarms.
    # pylint: disable=too-many-instance-attributes
    def __init__(self, options, maximum=True):
        self.namespace = options["namespace"]
        self.metric_name = options["metric_name"]
        self.hostclass = options["hostclass"]
        self.environment = options["environment"]
        self.team = options["team"]

        self.duration = int(options["duration"])
        self.period = int(options["period"])
        self.statistic = options["statistic"]
        self.custom_metric = options["custom_metric"].lower() == "true"
        self.log_pattern_metric = is_truthy(options.get("log_pattern_metric", ""))

        if "autoscaling_group_name" in options:
            self.autoscaling_group_name = options["autoscaling_group_name"]
        else:
            self.autoscaling_group_name = None

        if "es_domain_name" in options:
            self.es_domain_name = options["es_domain_name"]
        else:
            self.es_domain_name = None

        if "es_client_id" in options:
            self.es_client_id = options["es_client_id"]
        else:
            self.es_client_id = None

        if options["level"] not in NOTIFICATION_LEVELS:
            raise AlarmConfigError(
                "Level for {0} must be one of {1}"
                .format(self.metric_name, NOTIFICATION_LEVELS)
            )
        self._level = options["level"]

        if maximum:
            logger.debug("threshold_max: %s", options["threshold_max"])
            self.threshold = int(options["threshold_max"])
            self.threshold_type = "max"
        else:
            logger.debug("threshold_min: %s", options["threshold_min"])
            self.threshold = int(options["threshold_min"])
            self.threshold_type = "min"

    @property
    def dimensions(self):
        """
        Alarm Metric Dimensions
        """

        # metrics created from logs do not have dimensions
        if self.log_pattern_metric:
            return {}

        # RDS dimensions
        # hostclass name is the DBInstanceIdentifier
        if self.namespace == 'AWS/RDS':
            key = "DBInstanceIdentifier"
            # Format is {vpc}-{DbInstanceId} Underscore is not allowed in RDS instance name
            value = "-".join([self.environment, self.hostclass])
            return {key: value}

        if self.namespace == 'AWS/ELB':
            return {'LoadBalancerName': DiscoELB.get_elb_id(self.environment, self.hostclass)}

        if self.namespace == 'AWS/ES':
            return {
                'DomainName': self.es_domain_name,
                'ClientId': self.es_client_id
            }

        if self.custom_metric:
            key = "env_hostclass"
            value = "_".join([self.environment, self.hostclass])
        else:
            key = "AutoScalingGroupName"
            value = self.autoscaling_group_name
        return {key: value}

    @property
    def notification_topic(self):
        """
        The name of the topic that the alarm should route to.
        """
        return "_".join([self.team, self.environment, self._level])

    @property
    def name(self):
        """
        Alarm name
        """
        return "_".join([self.team, self.environment, self.hostclass, self.metric_name, self.threshold_type])

    @staticmethod
    def decode_alarm_name(name):
        """
        Decodes an alarm name into a dictionary with team, env, hostclass, metric_name and
        threshold_type.

        Raises AlarmConfigError if unable to parse the alarm name.
        """
        parts = name.split("_")
        if len(parts) >= 5:
            # alarm names that have extra underscores
            # Assume this means the metric name portion has underscores
            # FIXME this won't work if the alarm is missing a team name AND the metric name has underscores
            # FIXME consider using different delimiters in the future
            return {
                "team": parts[0],
                "env": parts[1],
                "hostclass": parts[2],
                # some metric names have '_'. Assume that any extra '_' in the alarm name are the metric name
                "metric_name": '_'.join(parts[3:-1]),
                "threshold_type": parts[-1],
            }
        elif len(parts) == 4:
            # alarm names that don't contain a team name and a metric name without underscores
            return {
                "team": None,
                "env": parts[0],
                "hostclass": parts[1],
                "metric_name": parts[2],
                "threshold_type": parts[3],
            }
        return {}
        # raise AlarmConfigError("Failed to decode alarm name {0}".format(name))

    def __repr__(self):
        return "Alarm:{0}.{1}".format(self.name, self.threshold)

    def to_metric_alarm(self, policy_arn):
        """
        Returns a MetricAlarm for a given policy.
        """
        return MetricAlarm(
            alarm_actions=[policy_arn],
            ok_actions=[policy_arn],
            comparison='>' if self.threshold_type == "max" else "<",
            dimensions=self.dimensions,
            evaluation_periods=self.duration,
            metric=self.metric_name,
            name=self.name,
            namespace=self.namespace,
            period=self.period,
            statistic=self.statistic,
            threshold=self.threshold,
        )

class DiscoAlarmsConfig(object):
    """
    Represenation of all alarms in config file.
    """

    def __init__(self, environment, config_file=None, autoscale=None, elasticsearch=None):
        self.config = read_config(config_file=(config_file or DEFAULT_CONFIG_FILE), environment=environment)
        self.environment = environment
        self.elasticsearch = elasticsearch or None
        self._autoscale = autoscale or None  # laziliy intialized
        self._defaults = None

    @property
    def defaults(self):
        """
        get default options for whole disco_alarm_config
        """
        self._defaults = self._defaults or self.get_defaults()
        return self._defaults

    @property
    def autoscale(self):
        """
        returns a lazily created disco autoscaling object
        """
        if not self._autoscale:
            self._autoscale = DiscoAutoscale(environment_name=self.environment)
        return self._autoscale

    @staticmethod
    def _decode_section_name(section):
        """
        Extract special alarm parameters which are stored in section name
        """
        section_segments = section.split(".")
        # If a section has four or more segments, treat the middle segments as the metric name and
        # the ending segment as the hostclass. This is to support some Elasticsearch metrics like
        # 'ClusterStatus.red'
        if len(section_segments) >= 4:
            team = section_segments[0]
            namespace = section_segments[1]
            metric_name = '.'.join(section_segments[2:-1])
            section_hostclass = section_segments[-1]
        elif len(section_segments) == 3:
            team, namespace, metric_name = section_segments
            section_hostclass = None
        else:
            raise AlarmConfigError("Skipping non-alarm like config section {0}.".format(section))
        return team, namespace, metric_name, section_hostclass

    # Our alarms get a little complicated, so theres more than a few variables involved...
    # pylint: disable=too-many-locals
    def _get_alarm_specification_dict(self, team, namespace, metric_name, hostclass, hostclass_specific,
                                      autoscaling_group_name=None):
        """
        Return options for alarm, inheriting values from parent.
        """
        options = copy.copy(self.defaults)
        parent_section_name = ".".join([team, namespace, metric_name])
        child_section_name = ".".join([team, namespace, metric_name, hostclass])

        if hostclass_specific:
            if self.config.has_section(parent_section_name):
                parent_options = dict(self.config.items(parent_section_name))
                options.update(parent_options)
            child_options = dict(self.config.items(child_section_name))
            options.update(child_options)
        else:
            if self.config.has_section(child_section_name):
                # The alarm will be created with the more specific hostclass config
                return None
            parent_options = dict(self.config.items(parent_section_name))
            options.update(parent_options)

        options["team"] = team
        options["namespace"] = namespace
        options["metric_name"] = metric_name
        options["hostclass"] = hostclass
        options["environment"] = self.environment

        # If an autoscaling group name was provided, include it
        if autoscaling_group_name:
            options["autoscaling_group_name"] = autoscaling_group_name
        # Otherwise, if using the EC2 namespace, we should lookup the autoscaling group name of the hostclass
        # and use that.
        elif namespace == "AWS/EC2":
            group = self.autoscale.get_existing_group(hostclass=hostclass)
            options["autoscaling_group_name"] = group.name

        if namespace == "AWS/ES" and self.elasticsearch:
            domain_name = self.elasticsearch.get_domain_name(hostclass)
            options["es_domain_name"] = domain_name
            options["es_client_id"] = self.elasticsearch.get_client_id(domain_name)

        # metrics for log files have special environment specific namespaces and hostclass specific names
        if is_truthy(options.get("log_pattern_metric")):
            options["namespace"] = namespace + '/' + self.environment
            options['metric_name'] = hostclass + '-' + metric_name

        return options

    def get_alarms(self, hostclass, autoscaling_group_name=None):
        """
        Returns list of DiscoAlarm objects for all the alarms associated with a hostclass.
        """
        alarms = []
        for section in self.config.sections():
            try:
                team, namespace, metric_name, section_hostclass = DiscoAlarmsConfig._decode_section_name(
                    section
                )
            except AlarmConfigError:
                if section in [NOTIFICATION_SECTION_NAME, DEFAULT_SECTION_NAME]:
                    continue
                else:
                    raise
            if section_hostclass and section_hostclass != hostclass:
                # Hostclass specific section for different hostclass
                continue

            if namespace == "AWS/ES" and not self.elasticsearch:
                logger.info("Skipping creation of %s.%s.%s.%s because no Elasticsearch object was provided",
                            team, namespace, metric_name, section_hostclass)
                continue

            options = self._get_alarm_specification_dict(
                team, namespace, metric_name, hostclass, section_hostclass == hostclass,
                autoscaling_group_name=autoscaling_group_name
            )

            if not options:
                # No alarms options means we don't need to make the alarm
                continue

            for threshold in ["threshold_max", "threshold_min"]:
                if threshold in options:
                    if options[threshold].isdigit():
                        alarms.append(
                            DiscoAlarmConfig(options, maximum=threshold == "threshold_max")
                        )
                    else:
                        raise AlarmConfigError(
                            "Not a valid threshold value for {0}: {1}".format(threshold, options))
        return alarms

    def get_notifications(self):
        """
        Returns list of DiscoNotification objects for all notifications in the current environment
        for the current pagerduty group.
        """
        notifications = []
        options = dict(self.config.items(NOTIFICATION_SECTION_NAME))

        for notification_name, endpoint_str in options.iteritems():
            notification_parts = notification_name.split("_")
            if len(notification_parts) < 3:
                raise AlarmConfigError("Underscores must separate team name, "
                                       "env name and notification level: %s" % notification_name)
            level = notification_parts[-1]
            env = notification_parts[1]

            if env != self.environment:
                continue
            if level not in NOTIFICATION_LEVELS:
                raise AlarmConfigError("Unsupported notification level: %s" % level)

            endpoints = endpoint_str.split(",")
            notifications.append(DiscoNotification(notification_name, endpoints))

        return notifications

    def get_defaults(self):
        """
        return dictionary of default values that will be used
        """
        return dict(self.config.items(DEFAULT_SECTION_NAME)) or {}
<|endoftext|>