<commit_before>import os, json
import pandas as pd
from jsonschema import validate, Draft4Validator
from igf_data.illumina.samplesheet import SampleSheet
from igf_data.utils.gviz_utils import convert_to_gviz_json_for_display

class Validate_project_and_samplesheet_metadata:
  '''
  A package for running validation checks for project and samplesheet metadata file

  :param samplesheet_file: A samplesheet input file
  :param metadata_files: A list of metadata input file
  :param samplesheet_schema: A json schema for samplesheet file validation
  :param metadata_schema: A json schema for metadata file validation
  '''
  def __init__(self,samplesheet_file,metadata_files,samplesheet_schema,metadata_schema):
    self.samplesheet_file=samplesheet_file
    self.metadata_files=metadata_files
    self.samplesheet_schema=samplesheet_schema
    self.metadata_schema=metadata_schema

  def get_samplesheet_validation_report(self):
    '''
    A method for running validation checks on input samplesheet file
    :returns: A list of errors or an empty list
    '''
    try:
      samplesheet=SampleSheet(infile=self.samplesheet_file)
      with open(self.samplesheet_schema,'r') as jp:
        json_data=json.load(jp)

      samplesheet_json_fields=list(json_data['items']['properties'].keys())
      errors=list()
      for header_name in samplesheet._data_header:
        if header_name not in samplesheet_json_fields:
          errors.append({'column':'',
                         'line':'',
                         'filename':os.path.basename(self.samplesheet_file),
                         'error':'Header {0} is not supported. Validation incomplete.'.\
                         format(header_name)}
                       )
      if len(errors)>0:
          return errors                                                         # stop checking samplesheet

      samplesheet_errors=samplesheet.validate_samplesheet_data(schema_json=self.samplesheet_schema)
      if len(samplesheet_errors)>0:
        errors=[{'column':'',
                 'line':'',
                 'filename':os.path.basename(self.samplesheet_file),
                 'error':err}
                 if isinstance(err,str) else
                {'column':err.schema_path[2],
                 'line':err.path[0]+1,
                 'filename':os.path.basename(self.samplesheet_file),
                 'error':err.message}
                for err in samplesheet_errors]
      return errors
    except:
      raise

  def get_metadata_validation_report(self):
    '''
    A method for running validation check on input metdata files
    :returns: A list of errors or an empty list
    '''
    try:
      error_list=list()
      with open(self.metadata_schema,'r') as jf:
        schema=json.load(jf)
      metadata_validator=Draft4Validator(schema)
      metadata_json_fields=list(schema['items']['properties'].keys())

      for metadata_file in self.metadata_files:
        metadata=pd.read_csv(metadata_file)
        metadata_error_list=list()
        for header_name in metadata.columns:
          if not header_name.startswith('Unnamed') and \
             not header_name in metadata_json_fields:                           # ignore any column without a header
            metadata_error_list.append({'column':'',
                                        'line':'',
                                        'filename':os.path.basename(metadata_file),
                                        'error':'Header {0} is not supported. Validation incomplete.'.\
                                        format(header_name)}
                                      )
        if len(metadata_error_list)>0:
          error_list.extend(metadata_error_list)
          continue                                                              # skip validation check for metadata file

        if 'library_source' in metadata.columns and \
             'library_strategy' in metadata.columns and \
             'experiment_type' in metadata.columns:
          error_msg=self.validate_metadata_library_type(library_source=,
                                                        library_strategy=,
                                                        experiment_type=
                                                       )
        metadata=metadata.fillna("").applymap(lambda x: str(x))
        if 'taxon_id' in metadata.columns:
          metadata['taxon_id']=metadata['taxon_id'].astype(str)

        json_data=metadata.to_dict(orient='records')
        errors=sorted(metadata_validator.iter_errors(json_data), key=lambda e: e.path)
        errors=[{'column':'',
                 'line':'',
                 'filename':os.path.basename(metadata_file),
                 'error':err}
                 if isinstance(err,str) else
                {'column':err.schema_path[2],
                 'line':err.path[0]+1,
                 'filename':os.path.basename(metadata_file),
                 'error':err.message}
                for err in errors]
        error_list.extend(errors)
      return error_list
    except:
      raise

  def get_merged_errors(self):
    '''
    A method for running the validation checks on input samplesheet metadata and samplesheet files
    :returns: A list of errors or an empty list
    '''
    try:
      all_errors=list()
      all_errors.extend(self.get_samplesheet_validation_report())
      all_errors.extend(self.get_metadata_validation_report())
      return all_errors
    except:
      raise

  def dump_error_to_csv(self,output_csv):
    '''
    A method for dumping list or errors to a csv file
    :returns: output csv file path if any errors found, or else None
    '''
    try:
      all_errors=self.get_merged_errors()
      if len(all_errors)==0:
        return None
      else:
        all_errors=pd.DataFrame(all_errors)
        if os.path.exists(output_csv):
          raise IOError('Output file {0} already present'.format(output_csv))

        all_errors.to_csv(output_csv,index=False)
        return output_csv
    except:
      raise

  def convert_errors_to_gviz(self,output_json=None):
    '''
    A method for converting the list of errors to gviz format json

    :param output_json: A output json file for saving data, default None
    :returns: A gviz json data block for the html output if output_json is None,
              or else None
    '''
    try:
      all_errors=self.get_merged_errors()
      if len(all_errors)==0:
        description={'error': ('string', 'Error')}
        columns_order=['error']
        data=[{'error':'No error found'}]
        json_data=convert_to_gviz_json_for_display(\
                    description=description,
                    data=data,
                    columns_order=columns_order
                  )
      else:
        description={'filename':('string', 'File name'),
                     'column':('string', 'Column name'),
                     'line' : ('string', 'Line No.'),
                     'error': ('string', 'Error')
                    }
        columns_order=['filename','column','line','error']
        json_data=convert_to_gviz_json_for_display(description=description,
                                                   data=all_errors,
                                                   columns_order=columns_order
                                                  )
      if output_json is None:
        return json_data
      else:
        with open(output_json,'w') as jf:
          jf.write(json_data)
        return None
    except:
      raise

  @staticmethod
  def check_metadata_library_by_row(data):
    '''
    A static method for checking library type metadata per row

    :param data: A pandas data series containing sample metadata
    :returns: An error message or None
    '''
    try:
      err=None
      if 'sample_igf_id' not in data:
        err='Sample igf id not found'
      else:
        if 'library_source' in data and \
           'library_strategy' in data and \
           'experiment_type' in data:
          err=Validate_project_and_samplesheet_metadata.validate_metadata_library_type(\
                sample_id=data['sample_igf_id'],
                library_source=data['library_source'],
                library_strategy=data['library_strategy'],
                experiment_type=data['experiment_type'])

      return err
    except:
      raise

  @staticmethod
  def validate_metadata_library_type(sample_id,library_source,library_strategy,experiment_type):
    '''
    A staticmethod for validating library metadata information for sample

    :param sample_id: Sample name
    :param library_source: Library source information
    :param library_strategy: Library strategy information
    :param experiment_type: Experiment type information
    :returns: A error message string or None
    '''
    try:
      error_msg=None
      if library_source == 'GENOMIC':
        if library_strategy not in ['WGS',
                                    'EXOME',
                                    'CHIP-SEQ',
                                    'ATAC-SEQ',
                                    'UNKNOWN'] or \
            experiment_type not in ['WGS',
                                    'EXOME',
                                    'ATAC-SEQ',
                                    'H3K4ME3',
                                    'H3K27ME3',
                                    'H3K27AC',
                                    'H3K36ME3',
                                    'HISTONE-NARROW',
                                    'HISTONE-BROAD',
                                    'TF',
                                    'UNKNOWN']:
          error_msg='{0}: library_strategy {1} or experiment_type {2} is not compatible with library_source {3}'.\
                    format(sample_id,library_strategy,experiment_type,library_source)
      elif library_source == 'TRANSCRIPTOMIC':
        if library_strategy not in ['RNA-SEQ'] or \
           experiment_type not in ['POLYA-RNA',
                                   'TOTAL-RNA',
                                   'SMALL-RNA']:
          error_msg='{0}: library_strategy {1} or experiment_type {2} is not compatible with library_source {3}'.\
                    format(sample_id,library_strategy,experiment_type,library_source)
      elif library_source == 'GENOMIC_SINGLE_CELL':
        if library_strategy not in ['UNKNOWN'] or \
           experiment_type not in ['UNKNOWN']:
          error_msg='{0}: library_strategy {1} or experiment_type {2} is not compatible with library_source {3}'.\
                    format(sample_id,library_strategy,experiment_type,library_source)
      elif library_source == 'TRANSCRIPTOMIC_SINGLE_CELL':
        if library_strategy not in ['RNA-SEQ'] or \
           experiment_type not in ['TENX-TRANSCRIPTOME',
                                   'DROP-SEQ-TRANSCRIPTOME']:
          error_msg='{0}: library_strategy {1} or experiment_type {2} is not compatible with library_source {3}'.\
                    format(sample_id,library_strategy,experiment_type,library_source)

      return error_msg
    except:
      raise<commit_msg>checking library metadata type<commit_after>import os, json
import pandas as pd
from jsonschema import validate, Draft4Validator
from igf_data.illumina.samplesheet import SampleSheet
from igf_data.utils.gviz_utils import convert_to_gviz_json_for_display

class Validate_project_and_samplesheet_metadata:
  '''
  A package for running validation checks for project and samplesheet metadata file

  :param samplesheet_file: A samplesheet input file
  :param metadata_files: A list of metadata input file
  :param samplesheet_schema: A json schema for samplesheet file validation
  :param metadata_schema: A json schema for metadata file validation
  '''
  def __init__(self,samplesheet_file,metadata_files,samplesheet_schema,metadata_schema):
    self.samplesheet_file=samplesheet_file
    self.metadata_files=metadata_files
    self.samplesheet_schema=samplesheet_schema
    self.metadata_schema=metadata_schema

  def get_samplesheet_validation_report(self):
    '''
    A method for running validation checks on input samplesheet file
    :returns: A list of errors or an empty list
    '''
    try:
      samplesheet=SampleSheet(infile=self.samplesheet_file)
      with open(self.samplesheet_schema,'r') as jp:
        json_data=json.load(jp)

      samplesheet_json_fields=list(json_data['items']['properties'].keys())
      errors=list()
      for header_name in samplesheet._data_header:
        if header_name not in samplesheet_json_fields:
          errors.append({'column':'',
                         'line':'',
                         'filename':os.path.basename(self.samplesheet_file),
                         'error':'Header {0} is not supported. Validation incomplete.'.\
                         format(header_name)}
                       )
      if len(errors)>0:
          return errors                                                         # stop checking samplesheet

      samplesheet_errors=samplesheet.validate_samplesheet_data(schema_json=self.samplesheet_schema)
      if len(samplesheet_errors)>0:
        errors=[{'column':'',
                 'line':'',
                 'filename':os.path.basename(self.samplesheet_file),
                 'error':err}
                 if isinstance(err,str) else
                {'column':err.schema_path[2],
                 'line':err.path[0]+1,
                 'filename':os.path.basename(self.samplesheet_file),
                 'error':err.message}
                for err in samplesheet_errors]
      return errors
    except:
      raise

  def get_metadata_validation_report(self):
    '''
    A method for running validation check on input metdata files
    :returns: A list of errors or an empty list
    '''
    try:
      error_list=list()
      with open(self.metadata_schema,'r') as jf:
        schema=json.load(jf)
      metadata_validator=Draft4Validator(schema)
      metadata_json_fields=list(schema['items']['properties'].keys())

      for metadata_file in self.metadata_files:
        metadata=pd.read_csv(metadata_file)
        metadata_error_list=list()
        for header_name in metadata.columns:
          if not header_name.startswith('Unnamed') and \
             not header_name in metadata_json_fields:                           # ignore any column without a header
            metadata_error_list.append({'column':'',
                                        'line':'',
                                        'filename':os.path.basename(metadata_file),
                                        'error':'Header {0} is not supported. Validation incomplete.'.\
                                        format(header_name)}
                                      )
        if len(metadata_error_list)>0:
          error_list.extend(metadata_error_list)
          continue                                                              # skip validation check for metadata file

        library_errors=list()
        library_errors=(metadata.apply(lambda x: self.check_metadata_library_by_row(data=x),
                                       axis=1,
                                       result_type=None))                       # check metadata per row
        library_errors=[i for i in library_errors
                            if i is not None]                                   # filter library errors
        if len(library_errors)>0:
          error_list.extend(library_errors)                                     # add library errors to the list of final errors

        metadata=metadata.fillna("").applymap(lambda x: str(x))
        if 'taxon_id' in metadata.columns:
          metadata['taxon_id']=metadata['taxon_id'].astype(str)

        json_data=metadata.to_dict(orient='records')
        errors=sorted(metadata_validator.iter_errors(json_data), key=lambda e: e.path)
        errors=[{'column':'',
                 'line':'',
                 'filename':os.path.basename(metadata_file),
                 'error':err}
                 if isinstance(err,str) else
                {'column':err.schema_path[2],
                 'line':err.path[0]+1,
                 'filename':os.path.basename(metadata_file),
                 'error':err.message}
                for err in errors]
        error_list.extend(errors)
      return error_list
    except:
      raise

  def get_merged_errors(self):
    '''
    A method for running the validation checks on input samplesheet metadata and samplesheet files
    :returns: A list of errors or an empty list
    '''
    try:
      all_errors=list()
      all_errors.extend(self.get_samplesheet_validation_report())
      all_errors.extend(self.get_metadata_validation_report())
      return all_errors
    except:
      raise

  def dump_error_to_csv(self,output_csv):
    '''
    A method for dumping list or errors to a csv file
    :returns: output csv file path if any errors found, or else None
    '''
    try:
      all_errors=self.get_merged_errors()
      if len(all_errors)==0:
        return None
      else:
        all_errors=pd.DataFrame(all_errors)
        if os.path.exists(output_csv):
          raise IOError('Output file {0} already present'.format(output_csv))

        all_errors.to_csv(output_csv,index=False)
        return output_csv
    except:
      raise

  def convert_errors_to_gviz(self,output_json=None):
    '''
    A method for converting the list of errors to gviz format json

    :param output_json: A output json file for saving data, default None
    :returns: A gviz json data block for the html output if output_json is None,
              or else None
    '''
    try:
      all_errors=self.get_merged_errors()
      if len(all_errors)==0:
        description={'error': ('string', 'Error')}
        columns_order=['error']
        data=[{'error':'No error found'}]
        json_data=convert_to_gviz_json_for_display(\
                    description=description,
                    data=data,
                    columns_order=columns_order
                  )
      else:
        description={'filename':('string', 'File name'),
                     'column':('string', 'Column name'),
                     'line' : ('string', 'Line No.'),
                     'error': ('string', 'Error')
                    }
        columns_order=['filename','column','line','error']
        json_data=convert_to_gviz_json_for_display(description=description,
                                                   data=all_errors,
                                                   columns_order=columns_order
                                                  )
      if output_json is None:
        return json_data
      else:
        with open(output_json,'w') as jf:
          jf.write(json_data)
        return None
    except:
      raise

  @staticmethod
  def check_metadata_library_by_row(data):
    '''
    A static method for checking library type metadata per row

    :param data: A pandas data series containing sample metadata
    :returns: An error message or None
    '''
    try:
      err=None
      if 'sample_igf_id' not in data:
        err='Sample igf id not found'
      else:
        if 'library_source' in data and \
           'library_strategy' in data and \
           'experiment_type' in data:
          err=Validate_project_and_samplesheet_metadata.validate_metadata_library_type(\
                sample_id=data['sample_igf_id'],
                library_source=data['library_source'],
                library_strategy=data['library_strategy'],
                experiment_type=data['experiment_type'])

      return err
    except:
      raise

  @staticmethod
  def validate_metadata_library_type(sample_id,library_source,library_strategy,experiment_type):
    '''
    A staticmethod for validating library metadata information for sample

    :param sample_id: Sample name
    :param library_source: Library source information
    :param library_strategy: Library strategy information
    :param experiment_type: Experiment type information
    :returns: A error message string or None
    '''
    try:
      error_msg=None
      if library_source == 'GENOMIC':
        if library_strategy not in ['WGS',
                                    'EXOME',
                                    'CHIP-SEQ',
                                    'ATAC-SEQ',
                                    'UNKNOWN'] or \
            experiment_type not in ['WGS',
                                    'EXOME',
                                    'ATAC-SEQ',
                                    'H3K4ME3',
                                    'H3K27ME3',
                                    'H3K27AC',
                                    'H3K36ME3',
                                    'HISTONE-NARROW',
                                    'HISTONE-BROAD',
                                    'TF',
                                    'UNKNOWN']:
          error_msg='{0}: library_strategy {1} or experiment_type {2} is not compatible with library_source {3}'.\
                    format(sample_id,library_strategy,experiment_type,library_source)
      elif library_source == 'TRANSCRIPTOMIC':
        if library_strategy not in ['RNA-SEQ'] or \
           experiment_type not in ['POLYA-RNA',
                                   'TOTAL-RNA',
                                   'SMALL-RNA']:
          error_msg='{0}: library_strategy {1} or experiment_type {2} is not compatible with library_source {3}'.\
                    format(sample_id,library_strategy,experiment_type,library_source)
      elif library_source == 'GENOMIC_SINGLE_CELL':
        if library_strategy not in ['UNKNOWN'] or \
           experiment_type not in ['UNKNOWN']:
          error_msg='{0}: library_strategy {1} or experiment_type {2} is not compatible with library_source {3}'.\
                    format(sample_id,library_strategy,experiment_type,library_source)
      elif library_source == 'TRANSCRIPTOMIC_SINGLE_CELL':
        if library_strategy not in ['RNA-SEQ'] or \
           experiment_type not in ['TENX-TRANSCRIPTOME',
                                   'DROP-SEQ-TRANSCRIPTOME']:
          error_msg='{0}: library_strategy {1} or experiment_type {2} is not compatible with library_source {3}'.\
                    format(sample_id,library_strategy,experiment_type,library_source)

      return error_msg
    except:
      raise<|endoftext|>


<commit_before>import argparse, unittest
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx
class FET(object):
    def __init__(self, on_conductance=1., off_conductance=0.1, threshold_voltage=-1.):
        self.on_conductance=on_conductance
        self.off_conductance=off_conductance
        self.threshold_voltage=threshold_voltage
    def get_conductance(self,gate_voltage):
        if gate_voltage<=self.threshold_voltage:
            return self.on_conductance
        else:
            return self.off_conductance
class Resistor(object):
    def __init__(self):
        self.conductance=1
    def get_conductance(self,dummy):
        return self.conductance
class Network(object):
    """
    This class implements a grid network of 1 Ohm resistors
    and allows specifying of the voltage sources and ground nodes
    the solve method returns the voltages at all of the nodes and
    saves them to an internal variable
    """
    def __init__(self,network_rows, network_columns, component, ground_nodes=[-1], voltage_sources=np.array([[0,5]])):
        #network parameters
        self.network_rows=network_rows
        self.network_columns=network_columns
        self.network_size=self.network_rows*self.network_columns

        #set network parameters
        if all(x>=0 for x in ground_nodes):
            self.ground_nodes=ground_nodes
        else:
            self.ground_nodes=[self.network_size-1]

        self.voltage_sources=voltage_sources
        self.network=nx.grid_2d_graph(self.network_rows,self.network_columns)
        self.adjacency_matrix=self.make_adjacency_matrix()
        self.gate_voltage=0
        self.make_components(component)

        # to make the matrices necessary for the MNA matrix equation
        self.mna_G=self.make_G()
        self.mna_A=self.make_A()
        self.mna_z=self.make_z()
    def getG_trivial(self, n1,n2):
        """holds enough information to reconstruct r1,c1 to r2,c2 information which equates to physical position"""
        return 1
    def make_components(self,component):
        for n1,n2 in self.network.edges():
            self.network.edge[n1][n2]['component']=component()
            self.network.edge[n1][n2]['conductance']=self.network.edge[n1][n2]['component'].get_conductance(1)
    def make_adjacency_matrix(self):
        adjacency_matrix=nx.to_numpy_matrix(self.network,nodelist=sorted(self.network.nodes()))
        return adjacency_matrix
    def make_G(self):
        n=self.network_size
        G=np.array(nx.to_numpy_matrix(self.network,nodelist=sorted(self.network.nodes()),weight='conductance'))
        for i in range(n):
            for j in range(n):
                if i==j:
                    G[i,j]=-sum(G[i])
        return G
    def make_A(self):
        # define symbols
        gnd=self.ground_nodes
        Vsrc=self.voltage_sources
        G=self.mna_G
        n=self.network_size

        B=np.zeros((n,len(Vsrc)))
        for i in range(n):
                if i in Vsrc[:,0]:
                    B[i,list(Vsrc[:,0]).index(i)]=1
        D=np.zeros((len(Vsrc),len(Vsrc)))
        BTD=np.append(B.T,D,axis=1)
        A=np.append(np.append(G,B,axis=1),BTD,axis=0)
        A=np.delete(np.delete(A,gnd,0),gnd,1)
        return A
    def make_z(self):
        r=self.network_rows
        c=self.network_columns
        return np.append(np.zeros((r*c-len(self.ground_nodes),1)), self.voltage_sources[:,1][:,None], axis=0)
    def make_voltages(self):
        # inserts the ground voltages back into x
        x=self.mna_x
        for i in self.ground_nodes:
            x=np.insert(x,i,0,axis=0)
        #splits the source currents from x
        x=x[:-1]
        self.node_voltages=x
        i=0
        for node in sorted(self.network.nodes()):
            self.network.node[node]['voltage']=float(x[i])
            i+=1
    def get_voltages(self):
        d=[]
        for node,data in self.network.nodes(data=True):
            d.append(data['voltage'])
        return d

    def solve_mna(self):
        self.mna_x=np.linalg.solve(self.mna_A,self.mna_z)
        self.update_network()
        return self.get_voltages()
    def update_network(self):
        self.make_voltages()
        self.make_currents()
    def show_network(self):
        r=self.network_rows
        c=self.network_columns
        pos={}
        for i in range(r):
            for j in range(c):
                pos[(i,j)]=j,i
        edges,weights = zip(*nx.get_edge_attributes(self.network,'current').items())

        nodes,voltages = zip(*nx.get_node_attributes(self.network,'voltage').items())

        # graph=nx.draw_networkx(self.network, pos, width=2, nodelist=nodes, node_color=voltages,  cmap=plt.get_cmap('plasma'), edge_cmap=plt.get_cmap('plasma'), edgelist=edges, edge_color=weights, node_size=30, with_labels=False)
        nodes=nx.draw_networkx_nodes(self.network, pos, width=2, nodelist=nodes, node_color=voltages,  cmap=plt.get_cmap('winter'), node_size=30, with_labels=False)
        edges=nx.draw_networkx_edges(self.network, pos, width=2, edgelist=edges, edge_color=weights,  edge_cmap=plt.get_cmap('autumn'), node_size=30, with_labels=False)
        plt.colorbar(edges,orientation='horizontal',label="Current")
        plt.colorbar(nodes,label="Node Voltage")
        plt.show()
    def make_currents(self):
        for n1,n2 in self.network.edges():
            if n1==(3,3) or n2==(3,3):
                print(n1,n2)
            g = float(self.network.edge[n1][n2]['conductance'])
            dV = float(self.network.node[n1]['voltage']) - float(self.network.node[n2]['voltage'])
            # to include current directionality one would have to
            #replace the abs with some sort of node-node direction rules
            self.network.edge[n1][n2]['current']= abs(g * dV)
    def get_currents(self):
        d=[]
        for n1,n2,data in self.network.edges(data=True):
            d.append(float(data['current']))
        return d

class NetworkTest(unittest.TestCase):
    def setUp(self):
        self.net=Network(10,10,Resistor,ground_nodes=[9,39,69,99],voltage_sources=np.array([[0,5],[30,5],[60,5],[90,5]]))

    def test_presolve_conductance(self):
        for n1,n2 in self.net.network.edges():
            self.assertEqual(self.net.network.edge[n1][n2]["conductance"],1)
    def test_presolve_voltages(self):
        # self.net.solve_mna()
        # print(self.net.network.edge[(0,0)][(0,1)])
        # print(self.net.network.node[(0,0)])
        pass

    def test_show(self):
        self.net.solve_mna()
        self.net.show_network()

if __name__ == "__main__":
    suite = unittest.TestLoader().loadTestsFromTestCase(NetworkTest)
    unittest.TextTestRunner(verbosity=2).run(suite)

    # net=Network(3,2,Resistor,ground_nodes=[3],voltage_sources=np.array([[1,5]]))
    # net.solve_mna()
    # net.show_network()
<commit_msg>added command line interface for running<commit_after>#!/usr/bin/env python3
import argparse, unittest, textwrap
import numpy as np
import matplotlib.pyplot as plt
import networkx as nx
class FET(object):
    def __init__(self, on_conductance=1., off_conductance=0.1, threshold_voltage=-1.):
        self.on_conductance=on_conductance
        self.off_conductance=off_conductance
        self.threshold_voltage=threshold_voltage
    def get_conductance(self,gate_voltage):
        if gate_voltage<=self.threshold_voltage:
            return self.on_conductance
        else:
            return self.off_conductance
class Resistor(object):
    def __init__(self):
        self.conductance=1
    def get_conductance(self,dummy):
        return self.conductance
class Network(object):
    """
    This class implements a grid network of 1 Ohm resistors
    and allows specifying of the voltage sources and ground nodes
    the solve method returns the voltages at all of the nodes and
    saves them to an internal variable
    """
    def __init__(self,network_rows, network_columns, component, ground_nodes=[-1], voltage_sources=np.array([[0,5]])):
        #network parameters
        self.network_rows=network_rows
        self.network_columns=network_columns
        self.network_size=self.network_rows*self.network_columns

        #set network parameters
        if all(x>=0 for x in ground_nodes):
            self.ground_nodes=ground_nodes
        else:
            self.ground_nodes=[self.network_size-1]

        self.voltage_sources=voltage_sources
        self.network=nx.grid_2d_graph(self.network_rows,self.network_columns)
        self.adjacency_matrix=self.make_adjacency_matrix()
        self.gate_voltage=0
        self.make_components(component)

        # to make the matrices necessary for the MNA matrix equation
        self.mna_G=self.make_G()
        self.mna_A=self.make_A()
        self.mna_z=self.make_z()
    def getG_trivial(self, n1,n2):
        """holds enough information to reconstruct r1,c1 to r2,c2 information which equates to physical position"""
        return 1
    def make_components(self,component):
        for n1,n2 in self.network.edges():
            self.network.edge[n1][n2]['component']=component()
            self.network.edge[n1][n2]['conductance']=self.network.edge[n1][n2]['component'].get_conductance(1)
    def make_adjacency_matrix(self):
        adjacency_matrix=nx.to_numpy_matrix(self.network,nodelist=sorted(self.network.nodes()))
        return adjacency_matrix
    def make_G(self):
        n=self.network_size
        G=np.array(nx.to_numpy_matrix(self.network,nodelist=sorted(self.network.nodes()),weight='conductance'))
        for i in range(n):
            for j in range(n):
                if i==j:
                    G[i,j]=-sum(G[i])
        return G
    def make_A(self):
        # define symbols
        gnd=self.ground_nodes
        Vsrc=self.voltage_sources
        G=self.mna_G
        n=self.network_size

        B=np.zeros((n,len(Vsrc)))
        for i in range(n):
                if i in Vsrc[:,0]:
                    B[i,list(Vsrc[:,0]).index(i)]=1
        D=np.zeros((len(Vsrc),len(Vsrc)))
        BTD=np.append(B.T,D,axis=1)
        A=np.append(np.append(G,B,axis=1),BTD,axis=0)
        A=np.delete(np.delete(A,gnd,0),gnd,1)
        return A
    def make_z(self):
        r=self.network_rows
        c=self.network_columns
        return np.append(np.zeros((r*c-len(self.ground_nodes),1)), self.voltage_sources[:,1][:,None], axis=0)
    def make_voltages(self):
        # inserts the ground voltages back into x
        x=self.mna_x
        for i in self.ground_nodes:
            x=np.insert(x,i,0,axis=0)
        #splits the source currents from x
        x=x[:-1]
        self.node_voltages=x
        i=0
        for node in sorted(self.network.nodes()):
            self.network.node[node]['voltage']=float(x[i])
            i+=1
    def get_voltages(self):
        d=[]
        for node,data in self.network.nodes(data=True):
            d.append(data['voltage'])
        return d

    def solve_mna(self):
        self.mna_x=np.linalg.solve(self.mna_A,self.mna_z)
        self.update_network()
        return self.get_voltages()
    def update_network(self):
        self.make_voltages()
        self.make_currents()
    def show_network(self):
        r=self.network_rows
        c=self.network_columns
        pos={}
        for i in range(r):
            for j in range(c):
                pos[(i,j)]=j,i
        edges,weights = zip(*nx.get_edge_attributes(self.network,'current').items())

        nodes,voltages = zip(*nx.get_node_attributes(self.network,'voltage').items())

        # graph=nx.draw_networkx(self.network, pos, width=2, nodelist=nodes, node_color=voltages,  cmap=plt.get_cmap('plasma'), edge_cmap=plt.get_cmap('plasma'), edgelist=edges, edge_color=weights, node_size=30, with_labels=False)
        nodes=nx.draw_networkx_nodes(self.network, pos, width=2, nodelist=nodes, node_color=voltages,  cmap=plt.get_cmap('winter'), node_size=30, with_labels=False)
        edges=nx.draw_networkx_edges(self.network, pos, width=2, edgelist=edges, edge_color=weights,  edge_cmap=plt.get_cmap('autumn'), node_size=30, with_labels=False)
        plt.colorbar(edges,orientation='horizontal',label="Current")
        plt.colorbar(nodes,label="Node Voltage")
        plt.show()
    def make_currents(self):
        for n1,n2 in self.network.edges():
            g = float(self.network.edge[n1][n2]['conductance'])
            dV = float(self.network.node[n1]['voltage']) - float(self.network.node[n2]['voltage'])
            # to include current directionality one would have to
            #replace the abs with some sort of node-node direction rules
            self.network.edge[n1][n2]['current']= abs(g * dV)
    def get_currents(self):
        d=[]
        for n1,n2,data in self.network.edges(data=True):
            d.append(float(data['current']))
        return d

class NetworkTest(unittest.TestCase):
    def setUp(self):
        self.net=Network(10,10,Resistor,ground_nodes=[9,39,69,99],voltage_sources=np.array([[0,5],[30,5],[60,5],[90,5]]))

    def test_presolve_conductance(self):
        for n1,n2 in self.net.network.edges():
            self.assertEqual(self.net.network.edge[n1][n2]["conductance"],1)
    def test_presolve_voltages(self):
        # self.net.solve_mna()
        # print(self.net.network.edge[(0,0)][(0,1)])
        # print(self.net.network.node[(0,0)])
        pass

    def test_show(self):
        self.net.solve_mna()
        # self.net.show_network()

if __name__ == "__main__":
    parser = argparse.ArgumentParser( formatter_class=argparse.RawDescriptionHelpFormatter, description=textwrap.dedent("""\
            Description of the various functions available:
                plotall   -  plots all data for each chip on a seperate plot and store in subdir plots/
                test      -  simply loads datasets use -i for interactive and -f to force reload"""))
    parser.add_argument("rows",type=int)
    parser.add_argument("columns",type=int)
    parser.add_argument("-t", "--test", action="store_true")
    parser.add_argument("--show", action="store_true")
    args = parser.parse_args()
    if args.rows and args.columns:
        net=Network(args.rows,args.columns,Resistor,ground_nodes=[args.rows*args.columns],voltage_sources=np.array([[0,5]]))
        net.solve_mna()
        if args.show:
            net.show_network()
    if args.test:
        suite = unittest.TestLoader().loadTestsFromTestCase(NetworkTest)
        unittest.TextTestRunner(verbosity=3).run(suite)

    # net=Network(3,2,Resistor,ground_nodes=[3],voltage_sources=np.array([[1,5]]))
    # net.solve_mna()
    # net.show_network()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os

from datetime import datetime
import json

#directory = r'C:\Source\hq\commcare-hq\django-hq\brac'
#directory = r'C:\Source\hq\commcare-hq\django-hq\pathfinder'
#directory = r'C:\Source\hq\data\brac-chp-jul-24\touse2'
#directory = r'C:\Source\hq\data\brac-chw-jul-24\touse2'
#directory = r'C:\Source\hq\data\brac\CHP\commcare\input'
#directory = r'C:\Source\hq\data\brac\CHW\commcare\input'
directory = r'C:\Source\hq\data\pathfinder\commcare\input'

#domain = "BRAC"
domain = "Pathfinder"
files = os.listdir(directory)
for file in files:
    if "xml" in file:
        filename = (os.path.join(directory,file))
        file = open(filename, "rb")
        dir, short_filename = os.path.split(filename)
        new_dir = os.path.join(dir, "out")
        payload = file.read()
        file.close()
        if not os.path.exists(new_dir):
            os.makedirs(new_dir)
        new_filename = short_filename.replace(".xml", ".postexport")
        (mode, ino, dev, nlink, uid, gid, size, atime, mtime, ctime) = os.stat(filename)
        headers = { "content-type" : "text/xml",
                "content-length" : str(len(payload)),
                "time-received" : str(datetime.fromtimestamp(mtime)),
                "original-ip" : "192.168.7.211",
                "domain" : domain
               }
        fout = open(os.path.join(new_dir, new_filename), 'w')
        jsoned = json.dumps(headers)
        fout.write(jsoned)
        fout.write("\n\n")
        fout.write(payload)
        fout.close()
<commit_msg>slight tweaks to the supervisor export script.  allow you to specify what directory you want the new files created in, rather than assuming you want to make an "out" folder wherever you're looking.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os

from datetime import datetime
import json

directory = r'C:\Source\hq\data\supervisor_export_all_08-08-2009\export\BRAC\CHW\commcare\input'
out_directory = r'C:\Source\hq\data\supervisor_export_all_08-08-2009\brac-pf'

domain = "BRAC"
#domain = "Pathfinder"
files = os.listdir(directory)
count = 0
for file in files:
    if "xml" in file:
        filename = (os.path.join(directory,file))
        file = open(filename, "rb")
        dir, short_filename = os.path.split(filename)
        #new_dir = os.path.join(dir, "out")
        new_dir = out_directory
        payload = file.read()
        file.close()
        if not os.path.exists(new_dir):
            os.makedirs(new_dir)
        new_filename = short_filename.replace(".xml", ".postexport")
        (mode, ino, dev, nlink, uid, gid, size, atime, mtime, ctime) = os.stat(filename)
        headers = { "content-type" : "text/xml",
                "content-length" : str(len(payload)),
                "time-received" : str(datetime.fromtimestamp(mtime)),
                "original-ip" : "192.168.7.211",
                "domain" : domain
               }
        fout = open(os.path.join(new_dir, new_filename), 'w')
        jsoned = json.dumps(headers)
        fout.write(jsoned)
        fout.write("\n\n")
        fout.write(payload)
        fout.close()
        count = count + 1

print "exported %s forms from %s to %s" % (count, directory, out_directory)<|endoftext|>


<commit_before>
import logging
import copy

import numpy as np

from .utils import bounded_sum, similarity

class Model(object):

    def __init__(self, num_primitives, num_actions):

        self.logger = logging.getLogger(self.__class__.__name__)

        self.SIMILARITY_THRESHOLD = 0.9       # real, 0 < x < 1
        self.MAX_ENTRIES = 10 ** 4               # integer, somewhat large
        self.CLEANING_PERIOD = 10 ** 5           # integer, somewhat large
        self.UPDATE_RATE = 0.1                # real, 0 < x < 1
        self.STEP_DISCOUNT = 0.5              # real, 0 < x < 1
        self.TRACE_LENGTH = 10                 # integer, small
        self.TRACE_DECAY_RATE = 0.2           # real, 0 < x < 1

        self.debug = False
        self.filename_postfix = '_model.mat'

        self.clean_count = 0

        self.cause = [[]]

        self.cause.append(np.zeros((num_primitives, 2*self.MAX_ENTRIES)))
        self.cause.append(np.zeros((num_actions, 2*self.MAX_ENTRIES)))

        self.effect = copy.deepcopy(self.cause)
        self.history = copy.deepcopy(self.cause)

        self.count = np.zeros(2*self.MAX_ENTRIES)
        self.reward_map = np.zeros(2*self.MAX_ENTRIES)
        self.goal_map = np.zeros(2*self.MAX_ENTRIES)

        self.trace_index = np.ones(self.TRACE_LENGTH)
        self.trace_reward = np.zeros(self.TRACE_LENGTH)

        # Initializes dummy transitions
        self.last_entry = 0
        self.cause[1][0,0] = 1
        self.effect[1][0,0] = 1
        #self.countp[0]= eps

    def add_group(self):
        size = (1, np.size(self.cause[-1], 2))
        self.history.append(np.zeros(size))
        self.cause.append(np.zeros(size))
        self.effect.append(np.zeros(size))

    def add_feature(self, nth_group, has_dummy):
        """
        add_feature to the model
        """

        self.history[nth_group] = np.vstack((self.history[nth_group],  np.zeros((1, np.size(self.history[nth_group], 2)))))
        self.cause[nth_group]   = np.vstack((self.cause[nth_group],  np.zeros((1, np.size(self.cause[nth_group],  2)))))
        self.effect[nth_group]  = np.vstack((self.effect[nth_group], np.zeros((1, np.size(self.effect[nth_group], 2)))))

        # if dummy feature is still in place, removes it
        if has_dummy:
            self.history[nth_group]   = self.history[nth_group][1:]
            self.cause[nth_group]  = self.cause[nth_group][1:]
            self.effect[nth_group] = self.effect[nth_group][1:]

    def train(self, effect, history, cause, reward):
        """
        Takes in 'cause' and 'effect'
        to train the model of the environment.
        """

        # TODO: udpate comments to reflect new representation

        eps = np.finfo(np.double).eps
        num_groups = len(cause)

        # Checks to see whether the new entry is already in the library
        # TODO: make the similarity threshold a function of the count? This would
        # allow often-observed transitions to require a closer fit, and populate
        # the model space more densely in areas where it accumulates more
        # observations.
        transition_similarity = similarity(history, self.history, range(0,self.last_entry))

        # If the cause doesn't match, the transition doesn't match
        cause_group = None
        cause_feature = None
        for group in cause[1:]:
            feature = group.ravel().nonzero()
            if feature:
                cause_group = group
                cause_feature = feature

        match_indices = []
        if cause_group:
            transition_similarity = transition_similarity * cause_group[cause_feature, 0:self.last_entry]
            match_indices = ( transition_similarity > self.SIMILARITY_THRESHOLD).ravel().nonzero()

        #if context and cause are sufficiently different,
        #adds a new entry to the model library
        if not match_indices:
            self.last_entry += 1
            matching_transition_index = self.last_entry
            for index in range(1,num_groups):
                self.history[index][:, self.last_entry] = history[index]
                self.cause[index][:, self.last_entry] = cause[index]
                self.effect[index][:, self.last_entry] = effect[index]

            self.count[matching_transition_index] =  1
            current_update_rate = 1

        #otherwise increments a nearby entry
        else:
            matching_transition_index = np.argmax(transition_similarity)

            self.count[matching_transition_index] += 1

            # modifies cause and effect
            # making the update rate a function of count allows updates to occur
            # more rapidly when there is little past experience to contradict them
            current_update_rate = (1 - self.UPDATE_RATE) / self.count(matching_transition_index) + self.UPDATE_RATE
            current_update_rate = min(1, current_update_rate)

            for index in range(1,num_groups):
                self.effect[index][:, matching_transition_index] = self.effect[index][:, matching_transition_index] * \
                    (1 - current_update_rate) + effect[index] * current_update_rate
                # TODO: update cause as well?
                #         self.cause[k][:, matching_transition_index] =

        #Performs credit assignment on the trace
        # updates the transition trace
        self.trace_index = np.array([self.trace_index[1:], matching_transition_index])

        # updates the reward trace
        self.trace_reward = np.array([self.trace_reward[1:], reward])

        credit = self.trace_reward[0]
        for i in range(1, self.TRACE_LENGTH):
            credit = bounded_sum( credit, self.trace_reward[i] * (1 - self.TRACE_DECAY_RATE) ** i)

        # updates the reward associated with the last entry in the trace
        update_index = self.trace_index[0]
        max_step_size = credit - self.reward_map[update_index]

        # # without assigning credit to the trace
        # update_index = matching_transition_index
        # max_step_size = reward - self.reward_map[update_index]

        self.reward_map[update_index] += max_step_size * current_update_rate

        #########################
        self.clean_count += 1

        # cleans out the library
        if self.last_entry >= self.MAX_ENTRIES:
            self.clean_count = self.CLEANING_PERIOD + 1

        if self.clean_count > self.CLEANING_PERIOD:

            self.count[:self.last_entry] -=  1 / self.count[:self.last_entry]
            forget_indices = (self.count[:self.last_entry] < eps).ravel().nonzero()

            for index in range(1,num_groups):
                self.history[index] = np.delete(self.history[index], forget_indices, 1)
                self.cause[index] = np.delete(self.cause[index], forget_indices, 1)
                self.effect[index] = np.delete(self.effect[index], forget_indices, 1)

            self.count = np.delete(self.count, forget_indices, 1)
            self.reward_map = np.delete(self.reward_map, forget_indices, 1)

            self.clean_count = 0
            self.last_entry -= len(forget_indices)
            if self.last_entry < 0:
                self.last_entry = 0

            #debug
            self.logger.debug('Library cleaning out %s entries to %s entries' % (len(forget_indices), self.last_entry))

        # pads the library if it has shrunk too far
        if np.size(self.effect[1], 1) < self.MAX_ENTRIES * 1.1:

            for index in range(1,num_groups):
                size = (self.effect[index].shape[0], self.MAX_ENTRIES)
                self.history[index] = np.hstack((self.history[index], np.zeros(size)))
                self.cause[index]  = np.hstack((self.cause[index], np.zeros(size)))
                self.effect[index] = np.hstack((self.effect[index], np.zeros(size)))

            self.count = np.hstack((self.count, np.zeros(self.MAX_ENTRIES)))
            self.reward_map = np.hstack((self.reward_map, np.zeros(self.MAX_ENTRIES)))

    def display_n_best(self, N):
        """
        provides a visual representation of the Nth cause-effect pair
        """

        sorted_indices = np.argsort(self.count[0:self.last_entry+1])

        sort_index = sorted_indices[0:N]

        for index in sorted_indices:
            self.display_pair(index)

    def display_pair(self, N):
        """
        provides a visual representation of the Nth cause-effect pair
        """

        num_groups = len(self.cause)

        for index, group in enumerate(self.cause):
            #subplot(num_groups-1, 2, (k-1)*2-1)
            #bar(self.cause{k}(:,N))
            #     disp(['self.cause{' num2str(k) '}(' num2str(N) ')'])
            #     self.cause{k}(:,N)
            #axis ([0 size(self.cause{k},1)+1 0 1])
            #ylabel(['g' num2str(k) ])
            #    xlabel(['max of ' num2str(max(self.cause{k}(:,N))) ])
            if index == 1:
                #title('model.cause for N = %s' % N)
                pass

            if index == num_groups - 1:
                #xlabel('count = %s' % self.count[index])
                pass

        for index, group in enumerate(self.cause[1:]):
            #subplot(num_groups-2, 2, index*2)
            #bar(self.effect{k}(:,N))
            #     disp(['self.effect{' num2str(k) '}(' num2str(N) ')'])
            #     self.effect{k}(:,N)
            #axis ([0 size(self.effect{k},1)+1 0 1])
            #ylabel(['g' num2str(k) ])
            #    xlabel(['max of ' num2str(max(self.effect{k}(:,N))) ])
            if index == 1:
                #title('model.effect for N = %s' % N)
                pass

    def display(self, N):
        """
        provides a visual representation of the Nth cause-effect pair
        """

        num_groups = len(self.cause)
        #
        # figure(177)
        # for k = 2:num_groups,
        #     subplot(num_groups-1, 2, (k-1)*2-1)
        #     bar(self.cause{k}(:,N))
        #     axis ([0 size(self.cause{k},1)+1 0 1])
        #     ylabel(['grp ' num2str(k) ])
        #     if (k == 2)
        #         title(['self.cause for N = ' num2str(N) ])
        #     end
        #     if (k == num_groups)
        #         xlabel(['count = ' num2str(self.count(N)) ])
        #     end
        # end
        #
        # for k = 2:num_groups,
        #     subplot(num_groups-1, 2, (k-1)*2)
        #     bar(self.effect{k}(:,N))
        #     axis ([0 size(self.effect{k},1)+1 0 1])
        #     ylabel(['grp ' num2str(k) ])
        #     if (k == 2)
        #         title(['self.effect for N = ' num2str(N) ])
        #     end
        # end

        # Text display
        #disp(['================Transition pair ' num2str(N) ' ================='])
        for index in range(1, num_groups):

            nonz_cause = len(self.cause[index][:,N].ravel().nonzero())
            nonz_effect = len(self.effect[index][:,N].ravel().nonzero())
            if ( (nonz_cause > 0) or (nonz_effect > 0)):

            #     # debug
            #     if ( (nonz_cause < 4) && (nonz_effect < 4))
            #
            #    disp(['================Transition pair ' num2str(N) ' ================='])

                print('  -----Group %s -----' % index)
                for subindex in range(0,len(self.effect[index][:,N])):
                    if (self.cause[index][subindex,N].nonzero()) or (self.effect[index][subindex,N].nonzero()):
                        print('    Feature             %s: %s %s' % (subindex, self.cause[index][subindex,N],
                                                                     self.effect[index][subindex,N]))

<commit_msg>model: convert more 1xN 2D to 1D vectors<commit_after>
import logging
import copy

import numpy as np

from .utils import bounded_sum, similarity

class Model(object):

    def __init__(self, num_primitives, num_actions):

        self.logger = logging.getLogger(self.__class__.__name__)

        self.SIMILARITY_THRESHOLD = 0.9       # real, 0 < x < 1
        self.MAX_ENTRIES = 10 ** 4               # integer, somewhat large
        self.CLEANING_PERIOD = 10 ** 5           # integer, somewhat large
        self.UPDATE_RATE = 0.1                # real, 0 < x < 1
        self.STEP_DISCOUNT = 0.5              # real, 0 < x < 1
        self.TRACE_LENGTH = 10                 # integer, small
        self.TRACE_DECAY_RATE = 0.2           # real, 0 < x < 1

        self.debug = False
        self.filename_postfix = '_model.mat'

        self.clean_count = 0

        self.cause = [[]]

        self.cause.append(np.zeros((num_primitives, 2*self.MAX_ENTRIES)))
        self.cause.append(np.zeros((num_actions, 2*self.MAX_ENTRIES)))

        self.effect = copy.deepcopy(self.cause)
        self.history = copy.deepcopy(self.cause)

        self.count = np.zeros(2*self.MAX_ENTRIES)
        self.reward_map = np.zeros(2*self.MAX_ENTRIES)
        self.goal_map = np.zeros(2*self.MAX_ENTRIES)

        self.trace_index = np.ones(self.TRACE_LENGTH)
        self.trace_reward = np.zeros(self.TRACE_LENGTH)

        # Initializes dummy transitions
        self.last_entry = 0
        self.cause[1][0,0] = 1
        self.effect[1][0,0] = 1
        #self.countp[0]= eps

    def add_group(self):
        size = (1, np.size(self.cause[-1], 2))
        self.history.append(np.zeros(size))
        self.cause.append(np.zeros(size))
        self.effect.append(np.zeros(size))

    def add_feature(self, nth_group, has_dummy):
        """
        add_feature to the model
        """

        self.history[nth_group] = np.vstack((self.history[nth_group],  np.zeros(self.history[nth_group].shape[1])))
        self.cause[nth_group]   = np.vstack((self.cause[nth_group],  np.zeros(self.cause[nth_group].shape[1])))
        self.effect[nth_group]  = np.vstack((self.effect[nth_group], np.zeros(self.effect[nth_group].shape[1])))

        # if dummy feature is still in place, removes it
        if has_dummy:
            self.history[nth_group]   = self.history[nth_group][1:]
            self.cause[nth_group]  = self.cause[nth_group][1:]
            self.effect[nth_group] = self.effect[nth_group][1:]

    def train(self, effect, history, cause, reward):
        """
        Takes in 'cause' and 'effect'
        to train the model of the environment.
        """

        # TODO: udpate comments to reflect new representation

        eps = np.finfo(np.double).eps
        num_groups = len(cause)

        # Checks to see whether the new entry is already in the library
        # TODO: make the similarity threshold a function of the count? This would
        # allow often-observed transitions to require a closer fit, and populate
        # the model space more densely in areas where it accumulates more
        # observations.
        transition_similarity = similarity(history, self.history, range(0,self.last_entry))

        # If the cause doesn't match, the transition doesn't match
        cause_group = None
        cause_feature = None
        for group in cause[1:]:
            feature = group.ravel().nonzero()
            if feature:
                cause_group = group
                cause_feature = feature

        match_indices = []
        if cause_group:
            transition_similarity = transition_similarity * cause_group[cause_feature, 0:self.last_entry]
            match_indices = ( transition_similarity > self.SIMILARITY_THRESHOLD).ravel().nonzero()

        #if context and cause are sufficiently different,
        #adds a new entry to the model library
        if not match_indices:
            self.last_entry += 1
            matching_transition_index = self.last_entry
            for index in range(1,num_groups):
                self.history[index][:, self.last_entry] = history[index]
                self.cause[index][:, self.last_entry] = cause[index]
                self.effect[index][:, self.last_entry] = effect[index]

            self.count[matching_transition_index] =  1
            current_update_rate = 1

        #otherwise increments a nearby entry
        else:
            matching_transition_index = np.argmax(transition_similarity)

            self.count[matching_transition_index] += 1

            # modifies cause and effect
            # making the update rate a function of count allows updates to occur
            # more rapidly when there is little past experience to contradict them
            current_update_rate = (1 - self.UPDATE_RATE) / self.count(matching_transition_index) + self.UPDATE_RATE
            current_update_rate = min(1, current_update_rate)

            for index in range(1,num_groups):
                self.effect[index][:, matching_transition_index] = self.effect[index][:, matching_transition_index] * \
                    (1 - current_update_rate) + effect[index] * current_update_rate
                # TODO: update cause as well?
                #         self.cause[k][:, matching_transition_index] =

        #Performs credit assignment on the trace
        # updates the transition trace
        self.trace_index = np.array([self.trace_index[1:], matching_transition_index])

        # updates the reward trace
        self.trace_reward = np.array([self.trace_reward[1:], reward])

        credit = self.trace_reward[0]
        for i in range(1, self.TRACE_LENGTH):
            credit = bounded_sum( credit, self.trace_reward[i] * (1 - self.TRACE_DECAY_RATE) ** i)

        # updates the reward associated with the last entry in the trace
        update_index = self.trace_index[0]
        max_step_size = credit - self.reward_map[update_index]

        # # without assigning credit to the trace
        # update_index = matching_transition_index
        # max_step_size = reward - self.reward_map[update_index]

        self.reward_map[update_index] += max_step_size * current_update_rate

        #########################
        self.clean_count += 1

        # cleans out the library
        if self.last_entry >= self.MAX_ENTRIES:
            self.clean_count = self.CLEANING_PERIOD + 1

        if self.clean_count > self.CLEANING_PERIOD:

            self.count[:self.last_entry] -=  1 / self.count[:self.last_entry]
            forget_indices = (self.count[:self.last_entry] < eps).ravel().nonzero()

            for index in range(1,num_groups):
                self.history[index] = np.delete(self.history[index], forget_indices, 1)
                self.cause[index] = np.delete(self.cause[index], forget_indices, 1)
                self.effect[index] = np.delete(self.effect[index], forget_indices, 1)

            self.count = np.delete(self.count, forget_indices, 1)
            self.reward_map = np.delete(self.reward_map, forget_indices, 1)

            self.clean_count = 0
            self.last_entry -= len(forget_indices)
            if self.last_entry < 0:
                self.last_entry = 0

            #debug
            self.logger.debug('Library cleaning out %s entries to %s entries' % (len(forget_indices), self.last_entry))

        # pads the library if it has shrunk too far
        if np.size(self.effect[1], 1) < self.MAX_ENTRIES * 1.1:

            for index in range(1,num_groups):
                size = (self.effect[index].shape[0], self.MAX_ENTRIES)
                self.history[index] = np.hstack((self.history[index], np.zeros(size)))
                self.cause[index]  = np.hstack((self.cause[index], np.zeros(size)))
                self.effect[index] = np.hstack((self.effect[index], np.zeros(size)))

            self.count = np.hstack((self.count, np.zeros(self.MAX_ENTRIES)))
            self.reward_map = np.hstack((self.reward_map, np.zeros(self.MAX_ENTRIES)))

    def display_n_best(self, N):
        """
        provides a visual representation of the Nth cause-effect pair
        """

        sorted_indices = np.argsort(self.count[0:self.last_entry+1])

        sort_index = sorted_indices[0:N]

        for index in sorted_indices:
            self.display_pair(index)

    def display_pair(self, N):
        """
        provides a visual representation of the Nth cause-effect pair
        """

        num_groups = len(self.cause)

        for index, group in enumerate(self.cause):
            #subplot(num_groups-1, 2, (k-1)*2-1)
            #bar(self.cause{k}(:,N))
            #     disp(['self.cause{' num2str(k) '}(' num2str(N) ')'])
            #     self.cause{k}(:,N)
            #axis ([0 size(self.cause{k},1)+1 0 1])
            #ylabel(['g' num2str(k) ])
            #    xlabel(['max of ' num2str(max(self.cause{k}(:,N))) ])
            if index == 1:
                #title('model.cause for N = %s' % N)
                pass

            if index == num_groups - 1:
                #xlabel('count = %s' % self.count[index])
                pass

        for index, group in enumerate(self.cause[1:]):
            #subplot(num_groups-2, 2, index*2)
            #bar(self.effect{k}(:,N))
            #     disp(['self.effect{' num2str(k) '}(' num2str(N) ')'])
            #     self.effect{k}(:,N)
            #axis ([0 size(self.effect{k},1)+1 0 1])
            #ylabel(['g' num2str(k) ])
            #    xlabel(['max of ' num2str(max(self.effect{k}(:,N))) ])
            if index == 1:
                #title('model.effect for N = %s' % N)
                pass

    def display(self, N):
        """
        provides a visual representation of the Nth cause-effect pair
        """

        num_groups = len(self.cause)
        #
        # figure(177)
        # for k = 2:num_groups,
        #     subplot(num_groups-1, 2, (k-1)*2-1)
        #     bar(self.cause{k}(:,N))
        #     axis ([0 size(self.cause{k},1)+1 0 1])
        #     ylabel(['grp ' num2str(k) ])
        #     if (k == 2)
        #         title(['self.cause for N = ' num2str(N) ])
        #     end
        #     if (k == num_groups)
        #         xlabel(['count = ' num2str(self.count(N)) ])
        #     end
        # end
        #
        # for k = 2:num_groups,
        #     subplot(num_groups-1, 2, (k-1)*2)
        #     bar(self.effect{k}(:,N))
        #     axis ([0 size(self.effect{k},1)+1 0 1])
        #     ylabel(['grp ' num2str(k) ])
        #     if (k == 2)
        #         title(['self.effect for N = ' num2str(N) ])
        #     end
        # end

        # Text display
        #disp(['================Transition pair ' num2str(N) ' ================='])
        for index in range(1, num_groups):

            nonz_cause = len(self.cause[index][:,N].ravel().nonzero())
            nonz_effect = len(self.effect[index][:,N].ravel().nonzero())
            if ( (nonz_cause > 0) or (nonz_effect > 0)):

            #     # debug
            #     if ( (nonz_cause < 4) && (nonz_effect < 4))
            #
            #    disp(['================Transition pair ' num2str(N) ' ================='])

                print('  -----Group %s -----' % index)
                for subindex in range(0,len(self.effect[index][:,N])):
                    if (self.cause[index][subindex,N].nonzero()) or (self.effect[index][subindex,N].nonzero()):
                        print('    Feature             %s: %s %s' % (subindex, self.cause[index][subindex,N],
                                                                     self.effect[index][subindex,N]))

<|endoftext|>


<commit_before>from unittest import TestCase, main, mock

from backend.waldur.logic.requests import *

class RequestTestCase(TestCase):
    def assert_correct_response_format(self, response, response_type="text"):
        self.assertIn('data', response)
        self.assertIn('type', response)
        self.assertEqual(response_type, response['type'])

def mocked_query_get_services_0_names(method, data, endpoint):
    return create_get_services_response()

def mocked_query_get_services_1_name(method, data, endpoint):
    return create_get_services_response("test1")

def mocked_query_get_services_2_names(method, data, endpoint):
    return create_get_services_response("test1", "test2")

def create_get_services_response(*names):
    return [
        {
            'services': [
                {
                    'name': name
                } for name in names
            ]
        }
    ]

class TestGetServicesRequests(RequestTestCase):
    def setUp(self):
        self.get_services = GetServicesRequest()
        self.get_services.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_services_0_names)
    def test_get_services_0(self, mock):
        response = self.get_services.process()
        self.assert_correct_response_format(response)

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_services_1_name)
    def test_get_services_1(self, mock):
        response = self.get_services.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_services_2_names)
    def test_get_services_2(self, mock):
        response = self.get_services.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])
        self.assertIn("test2", response['data'])

def mocked_query_get_projects_0_names(method, data, endpoint):
    return create_get_projects_response()

def mocked_query_get_projects_1_name(method, data, endpoint):
    return create_get_projects_response(("org1", {"test1"}))

def mocked_query_get_projects_2_names(method, data, endpoint):
    return create_get_projects_response(("org2", {"test1", "test2"}))

def mocked_query_get_projects_2_organisations(method, data, endpoint):
    return create_get_projects_response(("org1", {"test1", "test2"}), ("org2", {"test3"}))

def create_get_projects_response(*pairs):
    return [
        {
            'name': org,
            'projects': [
                {
                    'name': name
                } for name in names
            ]
        } for org, names in pairs
    ]

class TestGetProjectsRequests(RequestTestCase):
    def setUp(self):
        self.get_projects = GetProjectsRequest()
        self.get_projects.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_projects_0_names)
    def test_get_projects_0(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_projects_1_name)
    def test_get_projects_1(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        self.assertIn("org1", response['data'])
        self.assertIn("test1", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_projects_2_names)
    def test_get_projects_2(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        self.assertIn("org2", response['data'])
        self.assertIn("test1", response['data'])
        self.assertIn("test2", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_projects_2_organisations)
    def test_get_projects_3(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        self.assertIn("org1", response['data'])
        self.assertIn("test1", response['data'])
        self.assertIn("test2", response['data'])
        self.assertIn("org2", response['data'])
        self.assertIn("test3", response['data'])

def mocked_query_get_vms_0_names(method, data, endpoint):
    return create_get_vms_response()

def mocked_query_get_vms_1_name(method, data, endpoint):
    return create_get_vms_response(("test1", ["1.1", "1.0"]))

def mocked_query_get_vms_2_names(method, data, endpoint):
    return create_get_vms_response(("test1", ["1.1"]), ("test2", ["1.2"]))

def create_get_vms_response(*pairs):
    return [
        {
            'name': name,
            'external_ips': external_ips
        } for name, external_ips in pairs
    ]

class TestGetVmsRequests(RequestTestCase):
    def setUp(self):
        self.get_vms = GetVmsRequest()
        self.get_vms.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_vms_0_names)
    def test_get_vms_0(self, mock):
        response = self.get_vms.process()
        self.assert_correct_response_format(response)

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_vms_1_name)
    def test_get_vms_1(self, mock):
        response = self.get_vms.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])
        self.assertIn("1.1", response['data'])
        self.assertIn("1.0", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_vms_2_names)
    def test_get_vms_2(self, mock):
        response = self.get_vms.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])
        self.assertIn("1.1", response['data'])
        self.assertIn("test2", response['data'])
        self.assertIn("1.2", response['data'])

def mocked_query_get_organisations_0_names(method, data, endpoint):
    return create_get_organisations_response()

def mocked_query_get_organisations_1_name(method, data, endpoint):
    return create_get_organisations_response("test1")

def mocked_query_get_organisations_2_names(method, data, endpoint):
    return create_get_organisations_response("test1", "test2")

def create_get_organisations_response(*names):
    return [
        {
            'name': name
        } for name in names
    ]

class TestGetOrganisationsRequests(RequestTestCase):
    def setUp(self):
        self.get_organisations = GetOrganisationsRequest()
        self.get_organisations.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_organisations_0_names)
    def test_get_organisations_0(self, mock):
        response = self.get_organisations.process()
        self.assert_correct_response_format(response)

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_organisations_1_name)
    def test_get_organisations_1(self, mock):
        response = self.get_organisations.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_organisations_2_names)
    def test_get_organisations_2(self, mock):
        response = self.get_organisations.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])
        self.assertIn("test2", response['data'])

def mocked_query_get_total_cost_graph_empty(method, data, endpoint):
    return []

class TestGetTotalCostGraphRequest(RequestTestCase):
    def setUp(self):
        self.get_graph = GetTotalCostGraphRequest()
        self.get_graph.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_total_cost_graph_empty)
    def test_get_organisations_0(self, mock):
        response = self.get_graph.process()
        self.assert_correct_response_format(response, "graph")

def mocked_query_get_org_ids_1_name(method, data, endpoint):
    return create_get_org_ids_response(("test1", "id1"))

def mocked_query_get_org_ids_2_names(method, data, endpoint):
    return create_get_org_ids_response(("test1", "id1"), ("test2", "id2"))

def create_get_org_ids_response(*pairs):
    return [
        {
            'name': name,
            'uuid': uuid
        } for name, uuid in pairs
    ]

class TestGetOrganisationsAndIdsRequests(TestCase):
    def setUp(self):
        self.get_org_ids = GetOrganisationsAndIdsRequest()
        self.get_org_ids.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_org_ids_1_name)
    def test_get_org_ids_1(self, mock):
        response = self.get_org_ids.process()
        self.assertIn('test1', response)
        self.assertIn('id1', response['test1'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_org_ids_2_names)
    def test_get_org_ids_2(self, mock):
        response = self.get_org_ids.process()
        self.assertIn('test1', response)
        self.assertIn('id1', response['test1'])
        self.assertIn('test2', response)
        self.assertIn('id2', response['test2'])

def returns_true(*args):
    return True

class TestInputRequestSingleQuestion(TestCase):
    QUESTION = "What's up?"
    ANSWERS = ['y', 'yes']
    BAD_END_MSG = "how dare you"

    def setUp(self):
        self.ir = InputRequest([
            ('test', QA(self.QUESTION, self.ANSWERS))
        ], bad_end_msg=self.BAD_END_MSG)

    def test_set_get_input(self):
        self.ir.set_input("y")
        self.assertEqual("y", self.ir.input)
        self.assertEqual("y", self.ir.get_input())
        self.assertEqual(None, self.ir.input)
        self.assertEqual(None, self.ir.get_input())

    def test_handle_question_no_input(self):
        res = self.ir.handle_question()
        self.assertEqual(res, self.QUESTION)

    @mock.patch('backend.waldur.logic.requests.InputRequest.process', side_effect=returns_true)
    def test_handle_question_correct_input(self, mock):
        self.ir.set_input("y")
        # should reach process method, because we have no more questions after this one
        reached_process = self.ir.handle_question()
        self.assertTrue(reached_process)

        self.assertIn('test', self.ir.questions)
        # _evaluate must have been called, questions dict values should now have answers instead of questions
        self.assertEqual(self.ir.questions['test'], 'y')

    def test_handle_question_incorrect_input(self):
        self.ir.set_input("n")
        res = self.ir.handle_question()
        self.assertEqual(res, self.BAD_END_MSG)

class TestInputRequestMultipleQuestion(TestCase):
    I1 = "t1"
    Q1 = "q1"
    A1 = ['1']

    I2 = "t2"
    Q2 = "q2"
    A2 = ['2']

    I3 = "t3"
    Q3 = "q3"
    A3 = ['3']

    def setUp(self):
        self.ir = InputRequest([
            (self.I1, QA(self.Q1, self.A1)),
            (self.I2, QA(self.Q2, self.A2)),
            (self.I3, QA(self.Q3, self.A3)),
        ])

    @mock.patch('backend.waldur.logic.requests.InputRequest.process', side_effect=returns_true)
    def test_handle_all_questions(self, mock):
        question1 = self.ir.handle_question()
        self.assertEqual(question1, self.Q1)

        self.ir.set_input('1')
        question2 = self.ir.handle_question()
        self.assertEqual(question2, self.Q2)

        self.ir.set_input('2')
        question3 = self.ir.handle_question()
        self.assertEqual(question3, self.Q3)

        self.ir.set_input('3')
        reached_process = self.ir.handle_question()
        self.assertTrue(reached_process)

        self.assertEqual(self.ir.questions[self.I1], '1')
        self.assertEqual(self.ir.questions[self.I2], '2')
        self.assertEqual(self.ir.questions[self.I3], '3')

if __name__ == '__main__':
    main()
<commit_msg>systematic use case test? for UC-3 #28<commit_after>from unittest import TestCase, main, mock

from backend.waldur.logic.requests import *

class RequestTestCase(TestCase):
    def assert_correct_response_format(self, response, response_type="text"):
        self.assertIn('data', response)
        self.assertIn('type', response)
        self.assertEqual(response_type, response['type'])

def mocked_query_get_services_0_names(method, data, endpoint):
    return create_get_services_response()

def mocked_query_get_services_1_name(method, data, endpoint):
    return create_get_services_response("test1")

def mocked_query_get_services_2_names(method, data, endpoint):
    return create_get_services_response("test1", "test2")

def create_get_services_response(*names):
    return [
        {
            'services': [
                {
                    'name': name
                } for name in names
            ]
        }
    ]

class TestGetServicesRequests(RequestTestCase):
    def setUp(self):
        self.get_services = GetServicesRequest()
        self.get_services.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_services_0_names)
    def test_get_services_0(self, mock):
        response = self.get_services.process()
        self.assert_correct_response_format(response)

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_services_1_name)
    def test_get_services_1(self, mock):
        response = self.get_services.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_services_2_names)
    def test_get_services_2(self, mock):
        response = self.get_services.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])
        self.assertIn("test2", response['data'])

def mocked_query_get_projects_0_names(method, data, endpoint):
    return create_get_projects_response()

def mocked_query_get_projects_1_name(method, data, endpoint):
    return create_get_projects_response(("org1", {"test1"}))

def mocked_query_get_projects_2_names(method, data, endpoint):
    return create_get_projects_response(("org2", {"test1", "test2"}))

def mocked_query_get_projects_2_organisations(method, data, endpoint):
    return create_get_projects_response(("org1", {"test1", "test2"}), ("org2", {"test3"}))

def mocked_query_get_projects_2_organisations_1_empty(method, data, endpoint):
    return create_get_projects_response(("org1", {"test1", "test2"}), ("org2", {}))

def mocked_query_use_case_3_regular_get_projects(method, data, endpoint):
    o_1 = "Waldur Chatbot testbed (LTAT.05.005)"
    o_1_p_1 = "Waldur Chatbot testbed"
    o_1_p_2 = "2nd project"
    o_2 = "Waldur Maie"
    o_2_p_1 = "W-M project"
    return create_get_projects_response((o_1, (o_1_p_1, o_1_p_2)), (o_2, {o_2_p_1}))

def mocked_query_use_case_3_alt_a_get_projects(method, data, endpoint):
    o_1 = "Waldur Chatbot testbed (LTAT.05.005)"
    o_1_p_1 = "Waldur Chatbot testbed"
    o_2 = "Waldur Maie"
    return create_get_projects_response((o_1, {o_1_p_1}), (o_2, {}))

def mocked_query_use_case_3_alt_b_get_projects(method, data, endpoint):
    o_1 = "Waldur Chatbot testbed (LTAT.05.005)"
    return create_get_projects_response((o_1, {}))

def create_get_projects_response(*pairs):
    return [
        {
            'name': org,
            'projects': [
                {
                    'name': name
                } for name in names
            ]
        } for org, names in pairs
    ]

class TestGetProjectsRequests(RequestTestCase):
    def setUp(self):
        self.get_projects = GetProjectsRequest()
        self.get_projects.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_projects_0_names)
    def test_get_projects_0(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_projects_1_name)
    def test_get_projects_1(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        self.assertIn("org1", response['data'])
        self.assertIn("test1", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_projects_2_names)
    def test_get_projects_2(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        self.assertIn("org2", response['data'])
        self.assertIn("test1", response['data'])
        self.assertIn("test2", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_projects_2_organisations)
    def test_get_projects_3(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        self.assertIn("org1", response['data'])
        self.assertIn("test1", response['data'])
        self.assertIn("test2", response['data'])
        self.assertIn("org2", response['data'])
        self.assertIn("test3", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_projects_2_organisations_1_empty)
    def test_get_projects_4(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        self.assertIn("org1", response['data'])
        self.assertIn("test1", response['data'])
        self.assertIn("test2", response['data'])
        self.assertNotIn("org2", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_use_case_3_regular_get_projects)
    def test_get_projects_4(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        correct_response = "You have 3 projects in total.\nOrganisation 'Waldur Chatbot testbed (LTAT.05.005)':\n    "
        correct_response += "Waldur Chatbot testbed\n    2nd project\nOrganisation 'Waldur Maie':\n    W-M project"
        self.assertEqual(correct_response, response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_use_case_3_alt_a_get_projects)
    def test_get_projects_4(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        correct_response = "You have 1 project in total.\nOrganisation 'Waldur Chatbot testbed (LTAT.05.005)':\n    "
        correct_response += "Waldur Chatbot testbed"
        self.assertEqual(correct_response, response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_use_case_3_alt_b_get_projects)
    def test_get_projects_4(self, mock):
        response = self.get_projects.process()
        self.assert_correct_response_format(response)
        self.assertEqual("You don't have any projects.", response['data'])

def mocked_query_get_vms_0_names(method, data, endpoint):
    return create_get_vms_response()

def mocked_query_get_vms_1_name(method, data, endpoint):
    return create_get_vms_response(("test1", ["1.1", "1.0"]))

def mocked_query_get_vms_2_names(method, data, endpoint):
    return create_get_vms_response(("test1", ["1.1"]), ("test2", ["1.2"]))

def create_get_vms_response(*pairs):
    return [
        {
            'name': name,
            'external_ips': external_ips
        } for name, external_ips in pairs
    ]

class TestGetVmsRequests(RequestTestCase):
    def setUp(self):
        self.get_vms = GetVmsRequest()
        self.get_vms.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_vms_0_names)
    def test_get_vms_0(self, mock):
        response = self.get_vms.process()
        self.assert_correct_response_format(response)

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_vms_1_name)
    def test_get_vms_1(self, mock):
        response = self.get_vms.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])
        self.assertIn("1.1", response['data'])
        self.assertIn("1.0", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_vms_2_names)
    def test_get_vms_2(self, mock):
        response = self.get_vms.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])
        self.assertIn("1.1", response['data'])
        self.assertIn("test2", response['data'])
        self.assertIn("1.2", response['data'])

def mocked_query_get_organisations_0_names(method, data, endpoint):
    return create_get_organisations_response()

def mocked_query_get_organisations_1_name(method, data, endpoint):
    return create_get_organisations_response("test1")

def mocked_query_get_organisations_2_names(method, data, endpoint):
    return create_get_organisations_response("test1", "test2")

def create_get_organisations_response(*names):
    return [
        {
            'name': name
        } for name in names
    ]

class TestGetOrganisationsRequests(RequestTestCase):
    def setUp(self):
        self.get_organisations = GetOrganisationsRequest()
        self.get_organisations.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_organisations_0_names)
    def test_get_organisations_0(self, mock):
        response = self.get_organisations.process()
        self.assert_correct_response_format(response)

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_organisations_1_name)
    def test_get_organisations_1(self, mock):
        response = self.get_organisations.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_organisations_2_names)
    def test_get_organisations_2(self, mock):
        response = self.get_organisations.process()
        self.assert_correct_response_format(response)
        self.assertIn("test1", response['data'])
        self.assertIn("test2", response['data'])

def mocked_query_get_total_cost_graph_empty(method, data, endpoint):
    return []

class TestGetTotalCostGraphRequest(RequestTestCase):
    def setUp(self):
        self.get_graph = GetTotalCostGraphRequest()
        self.get_graph.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_total_cost_graph_empty)
    def test_get_organisations_0(self, mock):
        response = self.get_graph.process()
        self.assert_correct_response_format(response, "graph")

def mocked_query_get_org_ids_1_name(method, data, endpoint):
    return create_get_org_ids_response(("test1", "id1"))

def mocked_query_get_org_ids_2_names(method, data, endpoint):
    return create_get_org_ids_response(("test1", "id1"), ("test2", "id2"))

def create_get_org_ids_response(*pairs):
    return [
        {
            'name': name,
            'uuid': uuid
        } for name, uuid in pairs
    ]

class TestGetOrganisationsAndIdsRequests(TestCase):
    def setUp(self):
        self.get_org_ids = GetOrganisationsAndIdsRequest()
        self.get_org_ids.set_token("asd")

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_org_ids_1_name)
    def test_get_org_ids_1(self, mock):
        response = self.get_org_ids.process()
        self.assertIn('test1', response)
        self.assertIn('id1', response['test1'])

    @mock.patch('common.request.WaldurConnection.query', side_effect=mocked_query_get_org_ids_2_names)
    def test_get_org_ids_2(self, mock):
        response = self.get_org_ids.process()
        self.assertIn('test1', response)
        self.assertIn('id1', response['test1'])
        self.assertIn('test2', response)
        self.assertIn('id2', response['test2'])

def returns_true(*args):
    return True

class TestInputRequestSingleQuestion(TestCase):
    QUESTION = "What's up?"
    ANSWERS = ['y', 'yes']
    BAD_END_MSG = "how dare you"

    def setUp(self):
        self.ir = InputRequest([
            ('test', QA(self.QUESTION, self.ANSWERS))
        ], bad_end_msg=self.BAD_END_MSG)

    def test_set_get_input(self):
        self.ir.set_input("y")
        self.assertEqual("y", self.ir.input)
        self.assertEqual("y", self.ir.get_input())
        self.assertEqual(None, self.ir.input)
        self.assertEqual(None, self.ir.get_input())

    def test_handle_question_no_input(self):
        res = self.ir.handle_question()
        self.assertEqual(res, self.QUESTION)

    @mock.patch('backend.waldur.logic.requests.InputRequest.process', side_effect=returns_true)
    def test_handle_question_correct_input(self, mock):
        self.ir.set_input("y")
        # should reach process method, because we have no more questions after this one
        reached_process = self.ir.handle_question()
        self.assertTrue(reached_process)

        self.assertIn('test', self.ir.questions)
        # _evaluate must have been called, questions dict values should now have answers instead of questions
        self.assertEqual(self.ir.questions['test'], 'y')

    def test_handle_question_incorrect_input(self):
        self.ir.set_input("n")
        res = self.ir.handle_question()
        self.assertEqual(res, self.BAD_END_MSG)

class TestInputRequestMultipleQuestion(TestCase):
    I1 = "t1"
    Q1 = "q1"
    A1 = ['1']

    I2 = "t2"
    Q2 = "q2"
    A2 = ['2']

    I3 = "t3"
    Q3 = "q3"
    A3 = ['3']

    def setUp(self):
        self.ir = InputRequest([
            (self.I1, QA(self.Q1, self.A1)),
            (self.I2, QA(self.Q2, self.A2)),
            (self.I3, QA(self.Q3, self.A3)),
        ])

    @mock.patch('backend.waldur.logic.requests.InputRequest.process', side_effect=returns_true)
    def test_handle_all_questions(self, mock):
        question1 = self.ir.handle_question()
        self.assertEqual(question1, self.Q1)

        self.ir.set_input('1')
        question2 = self.ir.handle_question()
        self.assertEqual(question2, self.Q2)

        self.ir.set_input('2')
        question3 = self.ir.handle_question()
        self.assertEqual(question3, self.Q3)

        self.ir.set_input('3')
        reached_process = self.ir.handle_question()
        self.assertTrue(reached_process)

        self.assertEqual(self.ir.questions[self.I1], '1')
        self.assertEqual(self.ir.questions[self.I2], '2')
        self.assertEqual(self.ir.questions[self.I3], '3')

if __name__ == '__main__':
    main()
<|endoftext|>