<commit_before># -*- coding: utf-8 -*-
"""
This is the application file for the tool eTraGo.

Define your connection parameters and power flow settings before executing the function etrago.

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation; either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

__copyright__ = "Flensburg University of Applied Sciences, Europa-Universität Flensburg, Centre for Sustainable Energy Systems, DLR-Institute for Networked Energy Systems"
__license__ = "GNU Affero General Public License Version 3 (AGPL-3.0)"
__author__ = "ulfmueller, lukasol, wolfbunke, mariusves, s3pp"

import numpy as np
from numpy import genfromtxt
np.random.seed()
import time

import os

if not 'READTHEDOCS' in os.environ:
    # Sphinx does not run this code.
    # Do not import internal packages directly
    from etrago.tools.io import NetworkScenario, results_to_oedb
    from etrago.tools.plot import (plot_line_loading, plot_stacked_gen,
                                     add_coordinates, curtailment, gen_dist,
                                     storage_distribution, storage_expansion)
    from etrago.tools.utilities import (oedb_session, load_shedding, data_manipulation_sh,
                                    results_to_csv, parallelisation, pf_post_lopf,
                                    loading_minimization, calc_line_losses, group_parallel_lines)
    from etrago.cluster.networkclustering import busmap_from_psql, cluster_on_extra_high_voltage, kmean_clustering

from etrago.cluster.disaggregation import (MiniSolverDisaggregation,
                                           UniformDisaggregation)

args = {# Setup and Configuration:
        'db': 'oedb', # db session
        'gridversion': None, # None for model_draft or Version number (e.g. v0.2.11) for grid schema
        'method': 'lopf', # lopf or pf
        'pf_post_lopf': False, # state whether you want to perform a pf after a lopf simulation
        'start_snapshot': 2005,
        'end_snapshot' : 2006,
        'scn_name': 'SH NEP 2035', # state which scenario you want to run: Status Quo, NEP 2035, eGo100
        'solver': 'gurobi', # glpk, cplex or gurobi
        # Export options:
        'lpfile': False, # state if and where you want to save pyomo's lp file: False or /path/tofolder
        'results': False, # state if and where you want to save results as csv: False or /path/tofolder
        'export': False, # state if you want to export the results back to the database
        # Settings:
        'storage_extendable':True, # state if you want storages to be installed at each node if necessary.
        'generator_noise':True, # state if you want to apply a small generator noise
        'reproduce_noise': False, # state if you want to use a predefined set of random noise for the given scenario. if so, provide path, e.g. 'noise_values.csv'
        'minimize_loading':False,
        # Clustering:
        'k_mean_clustering': 10, # state if you want to perform a k-means clustering on the given network. State False or the value k (e.g. 20).
        'network_clustering': False, # state if you want to perform a clustering of HV buses to EHV buses.
        'disaggregation': 'uniform', # or None or mini
        # Simplifications:
        'parallelisation':False, # state if you want to run snapshots parallely.
        'skip_snapshots':4,
        'line_grouping': False, # state if you want to group lines running between the same buses.
        'branch_capacity_factor': 0.7, # globally extend or lower branch capacities
        'load_shedding':False, # meet the demand at very high cost; for debugging purposes.
        'comments':None }

def etrago(args):
    """The etrago function works with following arguments:

    Parameters
    ----------

    db (str):
    	'oedb',
        Name of Database session setting stored in config.ini of oemof.db

    gridversion (str):
        'v0.2.11',
        Name of the data version number of oedb: state 'None' for
        model_draft (sand-box) or an explicit version number
        (e.g. 'v0.2.10') for the grid schema.

    method (str):
        'lopf',
        Choose between a non-linear power flow ('pf') or
        a linear optimal power flow ('lopf').

    pf_post_lopf (bool):
        False,
        Option to run a non-linear power flow (pf) directly after the
        linear optimal power flow (and thus the dispatch) has finished.

    start_snapshot (int):
    	1,
        Start hour of the scenario year to be calculated.

    end_snapshot (int) :
    	2,
        End hour of the scenario year to be calculated.

    scn_name (str):
    	'Status Quo',
	Choose your scenario. Currently, there are three different
	scenarios: 'Status Quo', 'NEP 2035', 'eGo100'. If you do not
	want to use the full German dataset, you can use the excerpt of
	Schleswig-Holstein by adding the acronym SH to the scenario
	name (e.g. 'SH Status Quo').

    solver (str):
        'glpk',
        Choose your preferred solver. Current options: 'glpk' (open-source),
        'cplex' or 'gurobi'.

    lpfile (obj):
        False,
        State if and where you want to save pyomo's lp file. Options:
        False or '/path/tofolder'.

    results (obj):
        False,
        State if and where you want to save results as csv files.Options:
        False or '/path/tofolder'.

    export (bool):
        False,
        State if you want to export the results of your calculation
        back to the database.

    storage_extendable (bool):
        True,
        Choose if you want to allow to install extendable storages
        (unlimited in size) at each grid node in order to meet the flexibility demand.

    generator_noise (bool):
        True,
        Choose if you want to apply a small random noise to the marginal
        costs of each generator in order to prevent an optima plateau.

    reproduce_noise (obj):
        False,
        State if you want to use a predefined set of random noise for
        the given scenario. If so, provide path to the csv file,
        e.g. 'noise_values.csv'.

    minimize_loading (bool):
        False,

    k_mean_clustering (bool):
        False,
        State if you want to apply a clustering of all network buses down to
        only 'k' buses. The weighting takes place considering generation and load
        at each node. If so, state the number of k you want to apply. Otherwise
        put False. This function doesn't work together with 'line_grouping = True'
	    or 'network_clustering = True'.

    network_clustering (bool):
        False,
        Choose if you want to cluster the full HV/EHV dataset down to only the EHV
        buses. In that case, all HV buses are assigned to their closest EHV sub-station,
        taking into account the shortest distance on power lines.

    parallelisation (bool):
        False,
        Choose if you want to calculate a certain number of snapshots in parallel. If
        yes, define the respective amount in the if-clause execution below. Otherwise
        state False here.

    line_grouping (bool):
        True,
        State if you want to group lines that connect the same two buses into one system.

    branch_capacity_factor (numeric):
        1,
        Add a factor here if you want to globally change line capacities (e.g. to "consider"
        an (n-1) criterion or for debugging purposes.

    load_shedding (bool):
        False,
        State here if you want to make use of the load shedding function which is helpful when
        debugging: a very expensive generator is set to each bus and meets the demand when regular
        generators cannot do so.

    comments (str):
        None

    Result:
    -------

    """

    session = oedb_session(args['db'])

    # additional arguments cfgpath, version, prefix
    if args['gridversion'] == None:
        args['ormcls_prefix'] = 'EgoGridPfHv'
    else:
        args['ormcls_prefix'] = 'EgoPfHv'

    scenario = NetworkScenario(session,
                               version=args['gridversion'],
                               prefix=args['ormcls_prefix'],
                               method=args['method'],
                               start_snapshot=args['start_snapshot'],
                               end_snapshot=args['end_snapshot'],
                               scn_name=args['scn_name'])

    network = scenario.build_network()

    # add coordinates
    network = add_coordinates(network)

    # TEMPORARY vague adjustment due to transformer bug in data processing
    if args['gridversion'] == 'v0.2.11':
        network.transformers.x=network.transformers.x*0.0001

    clustering = None

    if args['branch_capacity_factor']:
        network.lines.s_nom = network.lines.s_nom*args['branch_capacity_factor']
        network.transformers.s_nom = network.transformers.s_nom*args['branch_capacity_factor']

    if args['generator_noise']:
        # create or reproduce generator noise
        if not args['reproduce_noise'] == False:
            noise_values = genfromtxt('noise_values.csv', delimiter=',')
            # add random noise to all generator
            network.generators.marginal_cost = noise_values
        else:
            noise_values = network.generators.marginal_cost + abs(np.random.normal(0,0.001,len(network.generators.marginal_cost)))
            np.savetxt("noise_values.csv", noise_values, delimiter=",")
            noise_values = genfromtxt('noise_values.csv', delimiter=',')
            # add random noise to all generator
            network.generators.marginal_cost = noise_values

    if args['storage_extendable']:
        # set virtual storages to be extendable
        if network.storage_units.carrier[network.storage_units.carrier== 'extendable_storage'].any() == 'extendable_storage':
            network.storage_units.loc[network.storage_units.carrier=='extendable_storage','p_nom_extendable'] = True
        # set virtual storage costs with regards to snapshot length
            network.storage_units.capital_cost = (network.storage_units.capital_cost /
            (8760//(args['end_snapshot']-args['start_snapshot']+1)))

    # for SH scenario run do data preperation:
    if args['scn_name'] == 'SH Status Quo' or args['scn_name'] == 'SH NEP 2035':
        data_manipulation_sh(network)

    # grouping of parallel lines
    if args['line_grouping']:
        group_parallel_lines(network)

    #load shedding in order to hunt infeasibilities
    if args['load_shedding']:
    	load_shedding(network)

    # network clustering
    if args['network_clustering']:
        network.generators.control="PV"
        busmap = busmap_from_psql(network, session, scn_name=args['scn_name'])
        network = cluster_on_extra_high_voltage(network, busmap, with_time=True)

    # k-mean clustering
    if not args['k_mean_clustering'] == False:
        clustering = kmean_clustering(network, n_clusters=args['k_mean_clustering'])
        original_network = network.copy()
        network = clustering.network.copy()

    # Branch loading minimization
    if args['minimize_loading']:
        extra_functionality = loading_minimization
    else:
        extra_functionality=None

    if args['skip_snapshots']:
        network.snapshots=network.snapshots[::args['skip_snapshots']]
        network.snapshot_weightings=network.snapshot_weightings[::args['skip_snapshots']]*args['skip_snapshots']

    # parallisation
    if args['parallelisation']:
        parallelisation(network, start_snapshot=args['start_snapshot'], end_snapshot=args['end_snapshot'],group_size=1, solver_name=args['solver'], extra_functionality=extra_functionality)
    # start linear optimal powerflow calculations
    elif args['method'] == 'lopf':
        x = time.time()
        network.lopf(network.snapshots, solver_name=args['solver'], extra_functionality=extra_functionality)
        y = time.time()
        z = (y - x) / 60 # z is time for lopf in minutes
    # start non-linear powerflow simulation
    elif args['method'] == 'pf':
        network.pf(scenario.timeindex)
       # calc_line_losses(network)

    if args['pf_post_lopf']:
        pf_post_lopf(network, scenario)
        calc_line_losses(network)

       # provide storage installation costs
    if sum(network.storage_units.p_nom_opt) != 0:
        installed_storages = network.storage_units[ network.storage_units.p_nom_opt!=0]
        storage_costs = sum(installed_storages.capital_cost * installed_storages.p_nom_opt)
        print("Investment costs for all storages in selected snapshots [EUR]:",round(storage_costs,2))

    if clustering:
        disagg = args.get('disaggregation')
        if disagg:
            if disagg == 'mini':
                disaggregation = MiniSolverDisaggregation(original_network,
                                                          network,
                                                          clustering)
            elif disagg == 'uniform':
                disaggregation = UniformDisaggregation(original_network,
                                                       network,
                                                       clustering)

            else:
                raise Exception('Invalid disaggregation command: ' + disagg)

            disaggregation.execute(scenario, solver=args['solver'])
            network = original_network

    # write lpfile to path
    if not args['lpfile'] == False:
        network.model.write(args['lpfile'], io_options={'symbolic_solver_labels':
                                                     True})
    # write PyPSA results back to database
    if args['export']:
        results_to_oedb(session, network, args, 'hv')

    # write PyPSA results to csv to path
    if not args['results'] == False:
        results_to_csv(network, args['results'])

    # close session
    session.close()

    return network

# execute etrago function
network = etrago(args)

if __name__ == '__main__':
    # execute etrago function
    network = etrago(args)
    # plots
    # make a line loading plot
    plot_line_loading(network)
    # plot stacked sum of nominal power for each generator type and timestep
    plot_stacked_gen(network, resolution="MW")
    # plot to show extendable storages
    storage_distribution(network)
<commit_msg>Disable snapshot skipping<commit_after># -*- coding: utf-8 -*-
"""
This is the application file for the tool eTraGo.

Define your connection parameters and power flow settings before executing the function etrago.

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation; either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

__copyright__ = "Flensburg University of Applied Sciences, Europa-Universität Flensburg, Centre for Sustainable Energy Systems, DLR-Institute for Networked Energy Systems"
__license__ = "GNU Affero General Public License Version 3 (AGPL-3.0)"
__author__ = "ulfmueller, lukasol, wolfbunke, mariusves, s3pp"

import numpy as np
from numpy import genfromtxt
np.random.seed()
import time

import os

if not 'READTHEDOCS' in os.environ:
    # Sphinx does not run this code.
    # Do not import internal packages directly
    from etrago.tools.io import NetworkScenario, results_to_oedb
    from etrago.tools.plot import (plot_line_loading, plot_stacked_gen,
                                     add_coordinates, curtailment, gen_dist,
                                     storage_distribution, storage_expansion)
    from etrago.tools.utilities import (oedb_session, load_shedding, data_manipulation_sh,
                                    results_to_csv, parallelisation, pf_post_lopf,
                                    loading_minimization, calc_line_losses, group_parallel_lines)
    from etrago.cluster.networkclustering import busmap_from_psql, cluster_on_extra_high_voltage, kmean_clustering

from etrago.cluster.disaggregation import (MiniSolverDisaggregation,
                                           UniformDisaggregation)

args = {# Setup and Configuration:
        'db': 'oedb', # db session
        'gridversion': None, # None for model_draft or Version number (e.g. v0.2.11) for grid schema
        'method': 'lopf', # lopf or pf
        'pf_post_lopf': False, # state whether you want to perform a pf after a lopf simulation
        'start_snapshot': 2005,
        'end_snapshot' : 2006,
        'scn_name': 'SH NEP 2035', # state which scenario you want to run: Status Quo, NEP 2035, eGo100
        'solver': 'gurobi', # glpk, cplex or gurobi
        # Export options:
        'lpfile': False, # state if and where you want to save pyomo's lp file: False or /path/tofolder
        'results': False, # state if and where you want to save results as csv: False or /path/tofolder
        'export': False, # state if you want to export the results back to the database
        # Settings:
        'storage_extendable':True, # state if you want storages to be installed at each node if necessary.
        'generator_noise':True, # state if you want to apply a small generator noise
        'reproduce_noise': False, # state if you want to use a predefined set of random noise for the given scenario. if so, provide path, e.g. 'noise_values.csv'
        'minimize_loading':False,
        # Clustering:
        'k_mean_clustering': 10, # state if you want to perform a k-means clustering on the given network. State False or the value k (e.g. 20).
        'network_clustering': False, # state if you want to perform a clustering of HV buses to EHV buses.
        'disaggregation': 'uniform', # or None or mini
        # Simplifications:
        'parallelisation':False, # state if you want to run snapshots parallely.
        'skip_snapshots':False,
        'line_grouping': False, # state if you want to group lines running between the same buses.
        'branch_capacity_factor': 0.7, # globally extend or lower branch capacities
        'load_shedding':False, # meet the demand at very high cost; for debugging purposes.
        'comments':None }

def etrago(args):
    """The etrago function works with following arguments:

    Parameters
    ----------

    db (str):
    	'oedb',
        Name of Database session setting stored in config.ini of oemof.db

    gridversion (str):
        'v0.2.11',
        Name of the data version number of oedb: state 'None' for
        model_draft (sand-box) or an explicit version number
        (e.g. 'v0.2.10') for the grid schema.

    method (str):
        'lopf',
        Choose between a non-linear power flow ('pf') or
        a linear optimal power flow ('lopf').

    pf_post_lopf (bool):
        False,
        Option to run a non-linear power flow (pf) directly after the
        linear optimal power flow (and thus the dispatch) has finished.

    start_snapshot (int):
    	1,
        Start hour of the scenario year to be calculated.

    end_snapshot (int) :
    	2,
        End hour of the scenario year to be calculated.

    scn_name (str):
    	'Status Quo',
	Choose your scenario. Currently, there are three different
	scenarios: 'Status Quo', 'NEP 2035', 'eGo100'. If you do not
	want to use the full German dataset, you can use the excerpt of
	Schleswig-Holstein by adding the acronym SH to the scenario
	name (e.g. 'SH Status Quo').

    solver (str):
        'glpk',
        Choose your preferred solver. Current options: 'glpk' (open-source),
        'cplex' or 'gurobi'.

    lpfile (obj):
        False,
        State if and where you want to save pyomo's lp file. Options:
        False or '/path/tofolder'.

    results (obj):
        False,
        State if and where you want to save results as csv files.Options:
        False or '/path/tofolder'.

    export (bool):
        False,
        State if you want to export the results of your calculation
        back to the database.

    storage_extendable (bool):
        True,
        Choose if you want to allow to install extendable storages
        (unlimited in size) at each grid node in order to meet the flexibility demand.

    generator_noise (bool):
        True,
        Choose if you want to apply a small random noise to the marginal
        costs of each generator in order to prevent an optima plateau.

    reproduce_noise (obj):
        False,
        State if you want to use a predefined set of random noise for
        the given scenario. If so, provide path to the csv file,
        e.g. 'noise_values.csv'.

    minimize_loading (bool):
        False,

    k_mean_clustering (bool):
        False,
        State if you want to apply a clustering of all network buses down to
        only 'k' buses. The weighting takes place considering generation and load
        at each node. If so, state the number of k you want to apply. Otherwise
        put False. This function doesn't work together with 'line_grouping = True'
	    or 'network_clustering = True'.

    network_clustering (bool):
        False,
        Choose if you want to cluster the full HV/EHV dataset down to only the EHV
        buses. In that case, all HV buses are assigned to their closest EHV sub-station,
        taking into account the shortest distance on power lines.

    parallelisation (bool):
        False,
        Choose if you want to calculate a certain number of snapshots in parallel. If
        yes, define the respective amount in the if-clause execution below. Otherwise
        state False here.

    line_grouping (bool):
        True,
        State if you want to group lines that connect the same two buses into one system.

    branch_capacity_factor (numeric):
        1,
        Add a factor here if you want to globally change line capacities (e.g. to "consider"
        an (n-1) criterion or for debugging purposes.

    load_shedding (bool):
        False,
        State here if you want to make use of the load shedding function which is helpful when
        debugging: a very expensive generator is set to each bus and meets the demand when regular
        generators cannot do so.

    comments (str):
        None

    Result:
    -------

    """

    session = oedb_session(args['db'])

    # additional arguments cfgpath, version, prefix
    if args['gridversion'] == None:
        args['ormcls_prefix'] = 'EgoGridPfHv'
    else:
        args['ormcls_prefix'] = 'EgoPfHv'

    scenario = NetworkScenario(session,
                               version=args['gridversion'],
                               prefix=args['ormcls_prefix'],
                               method=args['method'],
                               start_snapshot=args['start_snapshot'],
                               end_snapshot=args['end_snapshot'],
                               scn_name=args['scn_name'])

    network = scenario.build_network()

    # add coordinates
    network = add_coordinates(network)

    # TEMPORARY vague adjustment due to transformer bug in data processing
    if args['gridversion'] == 'v0.2.11':
        network.transformers.x=network.transformers.x*0.0001

    clustering = None

    if args['branch_capacity_factor']:
        network.lines.s_nom = network.lines.s_nom*args['branch_capacity_factor']
        network.transformers.s_nom = network.transformers.s_nom*args['branch_capacity_factor']

    if args['generator_noise']:
        # create or reproduce generator noise
        if not args['reproduce_noise'] == False:
            noise_values = genfromtxt('noise_values.csv', delimiter=',')
            # add random noise to all generator
            network.generators.marginal_cost = noise_values
        else:
            noise_values = network.generators.marginal_cost + abs(np.random.normal(0,0.001,len(network.generators.marginal_cost)))
            np.savetxt("noise_values.csv", noise_values, delimiter=",")
            noise_values = genfromtxt('noise_values.csv', delimiter=',')
            # add random noise to all generator
            network.generators.marginal_cost = noise_values

    if args['storage_extendable']:
        # set virtual storages to be extendable
        if network.storage_units.carrier[network.storage_units.carrier== 'extendable_storage'].any() == 'extendable_storage':
            network.storage_units.loc[network.storage_units.carrier=='extendable_storage','p_nom_extendable'] = True
        # set virtual storage costs with regards to snapshot length
            network.storage_units.capital_cost = (network.storage_units.capital_cost /
            (8760//(args['end_snapshot']-args['start_snapshot']+1)))

    # for SH scenario run do data preperation:
    if args['scn_name'] == 'SH Status Quo' or args['scn_name'] == 'SH NEP 2035':
        data_manipulation_sh(network)

    # grouping of parallel lines
    if args['line_grouping']:
        group_parallel_lines(network)

    #load shedding in order to hunt infeasibilities
    if args['load_shedding']:
    	load_shedding(network)

    # network clustering
    if args['network_clustering']:
        network.generators.control="PV"
        busmap = busmap_from_psql(network, session, scn_name=args['scn_name'])
        network = cluster_on_extra_high_voltage(network, busmap, with_time=True)

    # k-mean clustering
    if not args['k_mean_clustering'] == False:
        clustering = kmean_clustering(network, n_clusters=args['k_mean_clustering'])
        original_network = network.copy()
        network = clustering.network.copy()

    # Branch loading minimization
    if args['minimize_loading']:
        extra_functionality = loading_minimization
    else:
        extra_functionality=None

    if args['skip_snapshots']:
        network.snapshots=network.snapshots[::args['skip_snapshots']]
        network.snapshot_weightings=network.snapshot_weightings[::args['skip_snapshots']]*args['skip_snapshots']

    # parallisation
    if args['parallelisation']:
        parallelisation(network, start_snapshot=args['start_snapshot'], end_snapshot=args['end_snapshot'],group_size=1, solver_name=args['solver'], extra_functionality=extra_functionality)
    # start linear optimal powerflow calculations
    elif args['method'] == 'lopf':
        x = time.time()
        network.lopf(network.snapshots, solver_name=args['solver'], extra_functionality=extra_functionality)
        y = time.time()
        z = (y - x) / 60 # z is time for lopf in minutes
    # start non-linear powerflow simulation
    elif args['method'] == 'pf':
        network.pf(scenario.timeindex)
       # calc_line_losses(network)

    if args['pf_post_lopf']:
        pf_post_lopf(network, scenario)
        calc_line_losses(network)

       # provide storage installation costs
    if sum(network.storage_units.p_nom_opt) != 0:
        installed_storages = network.storage_units[ network.storage_units.p_nom_opt!=0]
        storage_costs = sum(installed_storages.capital_cost * installed_storages.p_nom_opt)
        print("Investment costs for all storages in selected snapshots [EUR]:",round(storage_costs,2))

    if clustering:
        disagg = args.get('disaggregation')
        if disagg:
            if disagg == 'mini':
                disaggregation = MiniSolverDisaggregation(original_network,
                                                          network,
                                                          clustering)
            elif disagg == 'uniform':
                disaggregation = UniformDisaggregation(original_network,
                                                       network,
                                                       clustering)

            else:
                raise Exception('Invalid disaggregation command: ' + disagg)

            disaggregation.execute(scenario, solver=args['solver'])
            network = original_network

    # write lpfile to path
    if not args['lpfile'] == False:
        network.model.write(args['lpfile'], io_options={'symbolic_solver_labels':
                                                     True})
    # write PyPSA results back to database
    if args['export']:
        results_to_oedb(session, network, args, 'hv')

    # write PyPSA results to csv to path
    if not args['results'] == False:
        results_to_csv(network, args['results'])

    # close session
    session.close()

    return network

# execute etrago function
network = etrago(args)

if __name__ == '__main__':
    # execute etrago function
    network = etrago(args)
    # plots
    # make a line loading plot
    plot_line_loading(network)
    # plot stacked sum of nominal power for each generator type and timestep
    plot_stacked_gen(network, resolution="MW")
    # plot to show extendable storages
    storage_distribution(network)
<|endoftext|>


<commit_before># Copyright 2017 The Forseti Security Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Creates a GCE instance template for Forseti Security."""

def GenerateConfig(context):
    """Generate configuration."""

    FORSETI_HOME = '$USER_HOME/forseti-security'

    DOWNLOAD_FORSETI = (
        "git clone {src_path}.git".format(
            src_path=context.properties['src-path']))

    FORSETI_VERSION = (
        "git checkout {forseti_version}".format(
            forseti_version=context.properties['forseti-version']))

    CLOUDSQL_CONN_STRING = '{}:{}:{}'.format(
        context.env['project'],
        '$(ref.cloudsql-instance.region)',
        '$(ref.cloudsql-instance.name)')

    SCANNER_BUCKET = context.properties['scanner-bucket']
    FORSETI_DB_NAME = context.properties['database-name']
    SERVICE_ACCOUNT_SCOPES =  context.properties['service-account-scopes']
    FORSETI_SERVER_CONF = '{}/configs/forseti_conf_server.yaml'.format(FORSETI_HOME)

    EXPORT_INITIALIZE_VARS = (
        'export SQL_PORT={0}\n'
        'export SQL_INSTANCE_CONN_STRING="{1}"\n'
        'export FORSETI_DB_NAME="{2}"\n')
    EXPORT_INITIALIZE_VARS = EXPORT_INITIALIZE_VARS.format(
        context.properties['db-port'],
        CLOUDSQL_CONN_STRING,
        FORSETI_DB_NAME)

    EXPORT_FORSETI_VARS = (
        'export FORSETI_HOME={forseti_home}\n'
        'export FORSETI_SERVER_CONF={forseti_server_conf}\n'
        ).format(forseti_home=FORSETI_HOME,
                 forseti_server_conf=FORSETI_SERVER_CONF)

    RUN_FREQUENCY = context.properties['run-frequency']
    EXTERNAL_PROJECT_ACCESS_SCANNER_CRON_ENABLED = context.properties['external-project-access-scanner-cron-enabled']
    EXTERNAL_PROJECT_ACCESS_SCANNER_CRON_SCHEDULE = context.properties['external-project-access-scanner-cron-schedule']

    resources = []

    deployment_name_splitted = context.env['deployment'].split('-')
    deployment_name_splitted.insert(len(deployment_name_splitted)-1, 'vm')
    instance_name = '-'.join(deployment_name_splitted)

    resources.append({
        'name': instance_name,
        'type': 'compute.v1.instance',
        'properties': {
            'zone': context.properties['zone'],
            'machineType': (
                'https://www.googleapis.com/compute/v1/projects/{}'
                '/zones/{}/machineTypes/{}'.format(
                context.env['project'], context.properties['zone'],
                context.properties['instance-type'])),
            'disks': [{
                'deviceName': 'boot',
                'type': 'PERSISTENT',
                'boot': True,
                'autoDelete': True,
                'initializeParams': {
                    'sourceImage': (
                        'https://www.googleapis.com/compute/v1'
                        '/projects/{}/global/images/family/{}'.format(
                            context.properties['image-project'],
                            context.properties['image-family']
                        )
                    )
                }
            }],
            'networkInterfaces': [{
                'network': (
                    'https://www.googleapis.com/compute/v1/'
                    'projects/{}/global/networks/{}'.format(
                    context.properties['vpc-host-project-id'],
                    context.properties['vpc-host-network'])),
                'accessConfigs': [{
                    'name': 'External NAT',
                    'type': 'ONE_TO_ONE_NAT'
                }],
                'subnetwork': (
                    'https://www.googleapis.com/compute/v1/'
                    'projects/{}/regions/{}/subnetworks/{}'.format(
                        context.properties['vpc-host-project-id'],
                        context.properties['region'],
                        context.properties['vpc-host-subnetwork']))
            }],
            'serviceAccounts': [{
                'email': context.properties['service-account'],
                'scopes': SERVICE_ACCOUNT_SCOPES,
            }],
            'metadata': {
                'items': [{
                    'key': 'startup-script',
                    'value': """#!/bin/bash
exec > /tmp/deployment.log
exec 2>&1

# Ubuntu available packages refresh.
sudo apt-get update -y

# Install Google Cloud SDK
sudo apt-get --assume-yes install google-cloud-sdk

USER_HOME=/home/ubuntu

# Install fluentd if necessary.
FLUENTD=$(ls /usr/sbin/google-fluentd)
if [ -z "$FLUENTD" ]; then
      cd $USER_HOME
      curl -sSO https://dl.google.com/cloudagents/install-logging-agent.sh
      bash install-logging-agent.sh
fi

# Check whether Cloud SQL proxy is installed.
CLOUD_SQL_PROXY=$(which cloud_sql_proxy)
if [ -z "$CLOUD_SQL_PROXY" ]; then
        cd $USER_HOME
        wget https://dl.google.com/cloudsql/cloud_sql_proxy.{cloudsql_arch}
        sudo mv cloud_sql_proxy.{cloudsql_arch} /usr/local/bin/cloud_sql_proxy
        chmod +x /usr/local/bin/cloud_sql_proxy
fi

# Install Forseti Security.
cd $USER_HOME
rm -rf *forseti*

# Download Forseti source code
{download_forseti}
cd forseti-security
git fetch --all
{checkout_forseti_version}

# Forseti Host Setup
sudo apt-get install -y git unzip

# Forseti host dependencies
sudo apt-get install -y $(cat install/dependencies/apt_packages.txt | grep -v "#" | xargs)

# Forseti dependencies
pip install --upgrade pip==9.0.3
pip install -q --upgrade setuptools wheel
pip install -q --upgrade -r requirements.txt

# Setup Forseti logging
touch /var/log/forseti.log
chown ubuntu:root /var/log/forseti.log
cp {forseti_home}/configs/logging/fluentd/forseti.conf /etc/google-fluentd/config.d/forseti.conf
cp {forseti_home}/configs/logging/logrotate/forseti /etc/logrotate.d/forseti
chmod 644 /etc/logrotate.d/forseti
service google-fluentd restart
logrotate /etc/logrotate.conf

# Change the access level of configs/ rules/ and run_forseti.sh
chmod -R ug+rwx {forseti_home}/configs {forseti_home}/rules {forseti_home}/install/gcp/scripts/run_forseti.sh

# Install Forseti
python setup.py install

# Export variables required by initialize_forseti_services.sh.
{export_initialize_vars}

# Export variables required by run_forseti.sh
{export_forseti_vars}

# Store the variables in /etc/profile.d/forseti_environment.sh
# so all the users will have access to them
echo "echo '{export_forseti_vars}' >> /etc/profile.d/forseti_environment.sh" | sudo sh

# Download server configuration from GCS
gsutil cp gs://{scanner_bucket}/configs/forseti_conf_server.yaml {forseti_server_conf}
gsutil cp -r gs://{scanner_bucket}/rules {forseti_home}/

# Start Forseti service depends on vars defined above.
bash ./install/gcp/scripts/initialize_forseti_services.sh

echo "Starting services."
systemctl start cloudsqlproxy
sleep 5

echo "Attempting to update database schema, if necessary."
python $USER_HOME/forseti-security/install/gcp/upgrade_tools/db_migrator.py

systemctl start forseti
echo "Success! The Forseti API server has been started."

# Create a Forseti env script
FORSETI_ENV="$(cat <<EOF
#!/bin/bash

export PATH=$PATH:/usr/local/bin

# Forseti environment variables
export FORSETI_HOME=/home/ubuntu/forseti-security
export FORSETI_SERVER_CONF=$FORSETI_HOME/configs/forseti_conf_server.yaml
export SCANNER_BUCKET={scanner_bucket}
EOF
)"
echo "$FORSETI_ENV" > $USER_HOME/forseti_env.sh

USER=ubuntu

# Use flock to prevent rerun of the same cron job when the previous job is still running.
# If the lock file does not exist under the tmp directory, it will create the file and put a lock on top of the file.
# When the previous cron job is not finished and the new one is trying to run, it will attempt to acquire the lock
# to the lock file and fail because the file is already locked by the previous process.
# The -n flag in flock will fail the process right away when the process is not able to acquire the lock so we won't
# queue up the jobs.
# If the cron job failed the acquire lock on the process, it will log a warning message to syslog.

(echo "{run_frequency} (/usr/bin/flock -n /home/ubuntu/forseti-security/forseti_cron_runner.lock $FORSETI_HOME/install/gcp/scripts/run_forseti.sh || echo '[forseti-security] Warning: New Forseti cron job will not be started, because previous Forseti job is still running.') 2>&1 | logger") | crontab -u $USER -
echo "Added the run_forseti.sh to crontab under user $USER"

# Adding a second cron job to run the External Project Access Scanner separately if it is enabled
if [ {external_project_access_scanner_cron_enabled} ]; then
    (crontab -l ; echo "{external_project_access_scanner_cron_schedule} (/usr/bin/flock -n /home/ubuntu/forseti-security/external_project_access_scanner_cron_runner.lock $FORSETI_HOME/install/gcp/scripts/run_external_project_access_scanner.sh || echo '[forseti-security] Warning: New external project access scanner cron job will not be started, because previous external project access scanner job is still running.') 2>&1 | logger")| crontab -
    echo "Added the run_external_project_access_scanner.sh to crontab under user $USER"
fi

echo "Execution of startup script finished"
""".format(
    # Cloud SQL properties
    cloudsql_arch = context.properties['cloudsqlproxy-os-arch'],

    # Install Forseti.
    download_forseti=DOWNLOAD_FORSETI,

    # Checkout Forseti version.
    checkout_forseti_version=FORSETI_VERSION,

    # Set ownership for Forseti conf and rules dirs
    forseti_home=FORSETI_HOME,

    # Download the Forseti conf and rules.
    scanner_bucket=SCANNER_BUCKET,
    forseti_server_conf=FORSETI_SERVER_CONF,

    # Env variables for Explain
    export_initialize_vars=EXPORT_INITIALIZE_VARS,

    # Env variables for Forseti
    export_forseti_vars=EXPORT_FORSETI_VARS,

    # Forseti run frequency
    run_frequency=RUN_FREQUENCY,
    external_project_access_scanner_cron_schedule=EXTERNAL_PROJECT_ACCESS_SCANNER_CRON_SCHEDULE,
    external_project_access_scanner_cron_enabled=EXTERNAL_PROJECT_ACCESS_SCANNER_CRON_ENABLED,
)
                }]
            }
        }
    })
    return {'resources': resources}
<commit_msg>fix the cron job creation issue<commit_after># Copyright 2017 The Forseti Security Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Creates a GCE instance template for Forseti Security."""

def GenerateConfig(context):
    """Generate configuration."""

    FORSETI_HOME = '$USER_HOME/forseti-security'

    DOWNLOAD_FORSETI = (
        "git clone {src_path}.git".format(
            src_path=context.properties['src-path']))

    FORSETI_VERSION = (
        "git checkout {forseti_version}".format(
            forseti_version=context.properties['forseti-version']))

    CLOUDSQL_CONN_STRING = '{}:{}:{}'.format(
        context.env['project'],
        '$(ref.cloudsql-instance.region)',
        '$(ref.cloudsql-instance.name)')

    SCANNER_BUCKET = context.properties['scanner-bucket']
    FORSETI_DB_NAME = context.properties['database-name']
    SERVICE_ACCOUNT_SCOPES =  context.properties['service-account-scopes']
    FORSETI_SERVER_CONF = '{}/configs/forseti_conf_server.yaml'.format(FORSETI_HOME)

    EXPORT_INITIALIZE_VARS = (
        'export SQL_PORT={0}\n'
        'export SQL_INSTANCE_CONN_STRING="{1}"\n'
        'export FORSETI_DB_NAME="{2}"\n')
    EXPORT_INITIALIZE_VARS = EXPORT_INITIALIZE_VARS.format(
        context.properties['db-port'],
        CLOUDSQL_CONN_STRING,
        FORSETI_DB_NAME)

    EXPORT_FORSETI_VARS = (
        'export FORSETI_HOME={forseti_home}\n'
        'export FORSETI_SERVER_CONF={forseti_server_conf}\n'
        ).format(forseti_home=FORSETI_HOME,
                 forseti_server_conf=FORSETI_SERVER_CONF)

    RUN_FREQUENCY = context.properties['run-frequency']
    EXTERNAL_PROJECT_ACCESS_SCANNER_CRON_ENABLED = context.properties['external-project-access-scanner-cron-enabled']
    EXTERNAL_PROJECT_ACCESS_SCANNER_CRON_SCHEDULE = context.properties['external-project-access-scanner-cron-schedule']

    resources = []

    deployment_name_splitted = context.env['deployment'].split('-')
    deployment_name_splitted.insert(len(deployment_name_splitted)-1, 'vm')
    instance_name = '-'.join(deployment_name_splitted)

    resources.append({
        'name': instance_name,
        'type': 'compute.v1.instance',
        'properties': {
            'zone': context.properties['zone'],
            'machineType': (
                'https://www.googleapis.com/compute/v1/projects/{}'
                '/zones/{}/machineTypes/{}'.format(
                context.env['project'], context.properties['zone'],
                context.properties['instance-type'])),
            'disks': [{
                'deviceName': 'boot',
                'type': 'PERSISTENT',
                'boot': True,
                'autoDelete': True,
                'initializeParams': {
                    'sourceImage': (
                        'https://www.googleapis.com/compute/v1'
                        '/projects/{}/global/images/family/{}'.format(
                            context.properties['image-project'],
                            context.properties['image-family']
                        )
                    )
                }
            }],
            'networkInterfaces': [{
                'network': (
                    'https://www.googleapis.com/compute/v1/'
                    'projects/{}/global/networks/{}'.format(
                    context.properties['vpc-host-project-id'],
                    context.properties['vpc-host-network'])),
                'accessConfigs': [{
                    'name': 'External NAT',
                    'type': 'ONE_TO_ONE_NAT'
                }],
                'subnetwork': (
                    'https://www.googleapis.com/compute/v1/'
                    'projects/{}/regions/{}/subnetworks/{}'.format(
                        context.properties['vpc-host-project-id'],
                        context.properties['region'],
                        context.properties['vpc-host-subnetwork']))
            }],
            'serviceAccounts': [{
                'email': context.properties['service-account'],
                'scopes': SERVICE_ACCOUNT_SCOPES,
            }],
            'metadata': {
                'items': [{
                    'key': 'startup-script',
                    'value': """#!/bin/bash
exec > /tmp/deployment.log
exec 2>&1

# Ubuntu available packages refresh.
sudo apt-get update -y

# Install Google Cloud SDK
sudo apt-get --assume-yes install google-cloud-sdk

USER_HOME=/home/ubuntu

# Install fluentd if necessary.
FLUENTD=$(ls /usr/sbin/google-fluentd)
if [ -z "$FLUENTD" ]; then
      cd $USER_HOME
      curl -sSO https://dl.google.com/cloudagents/install-logging-agent.sh
      bash install-logging-agent.sh
fi

# Check whether Cloud SQL proxy is installed.
CLOUD_SQL_PROXY=$(which cloud_sql_proxy)
if [ -z "$CLOUD_SQL_PROXY" ]; then
        cd $USER_HOME
        wget https://dl.google.com/cloudsql/cloud_sql_proxy.{cloudsql_arch}
        sudo mv cloud_sql_proxy.{cloudsql_arch} /usr/local/bin/cloud_sql_proxy
        chmod +x /usr/local/bin/cloud_sql_proxy
fi

# Install Forseti Security.
cd $USER_HOME
rm -rf *forseti*

# Download Forseti source code
{download_forseti}
cd forseti-security
git fetch --all
{checkout_forseti_version}

# Forseti Host Setup
sudo apt-get install -y git unzip

# Forseti host dependencies
sudo apt-get install -y $(cat install/dependencies/apt_packages.txt | grep -v "#" | xargs)

# Forseti dependencies
pip install --upgrade pip==9.0.3
pip install -q --upgrade setuptools wheel
pip install -q --upgrade -r requirements.txt

# Setup Forseti logging
touch /var/log/forseti.log
chown ubuntu:root /var/log/forseti.log
cp {forseti_home}/configs/logging/fluentd/forseti.conf /etc/google-fluentd/config.d/forseti.conf
cp {forseti_home}/configs/logging/logrotate/forseti /etc/logrotate.d/forseti
chmod 644 /etc/logrotate.d/forseti
service google-fluentd restart
logrotate /etc/logrotate.conf

# Change the access level of configs/ rules/ and run_forseti.sh
chmod -R ug+rwx {forseti_home}/configs {forseti_home}/rules {forseti_home}/install/gcp/scripts/run_forseti.sh

# Install Forseti
python setup.py install

# Export variables required by initialize_forseti_services.sh.
{export_initialize_vars}

# Export variables required by run_forseti.sh
{export_forseti_vars}

# Store the variables in /etc/profile.d/forseti_environment.sh
# so all the users will have access to them
echo "echo '{export_forseti_vars}' >> /etc/profile.d/forseti_environment.sh" | sudo sh

# Download server configuration from GCS
gsutil cp gs://{scanner_bucket}/configs/forseti_conf_server.yaml {forseti_server_conf}
gsutil cp -r gs://{scanner_bucket}/rules {forseti_home}/

# Start Forseti service depends on vars defined above.
bash ./install/gcp/scripts/initialize_forseti_services.sh

echo "Starting services."
systemctl start cloudsqlproxy
sleep 5

echo "Attempting to update database schema, if necessary."
python $USER_HOME/forseti-security/install/gcp/upgrade_tools/db_migrator.py

systemctl start forseti
echo "Success! The Forseti API server has been started."

# Create a Forseti env script
FORSETI_ENV="$(cat <<EOF
#!/bin/bash

export PATH=$PATH:/usr/local/bin

# Forseti environment variables
export FORSETI_HOME=/home/ubuntu/forseti-security
export FORSETI_SERVER_CONF=$FORSETI_HOME/configs/forseti_conf_server.yaml
export SCANNER_BUCKET={scanner_bucket}
EOF
)"
echo "$FORSETI_ENV" > $USER_HOME/forseti_env.sh

USER=ubuntu

# Use flock to prevent rerun of the same cron job when the previous job is still running.
# If the lock file does not exist under the tmp directory, it will create the file and put a lock on top of the file.
# When the previous cron job is not finished and the new one is trying to run, it will attempt to acquire the lock
# to the lock file and fail because the file is already locked by the previous process.
# The -n flag in flock will fail the process right away when the process is not able to acquire the lock so we won't
# queue up the jobs.
# If the cron job failed the acquire lock on the process, it will log a warning message to syslog.

(echo "{run_frequency} (/usr/bin/flock -n /home/ubuntu/forseti-security/forseti_cron_runner.lock $FORSETI_HOME/install/gcp/scripts/run_forseti.sh || echo '[forseti-security] Warning: New Forseti cron job will not be started, because previous Forseti job is still running.') 2>&1 | logger") | crontab -u $USER -
echo "Added the run_forseti.sh to crontab under user $USER"

# Adding a second cron job to run the External Project Access Scanner separately if it is enabled
if [ {external_project_access_scanner_cron_enabled} ]; then
    (crontab -u $USER -l; echo "{external_project_access_scanner_cron_schedule} (/usr/bin/flock -n /home/ubuntu/forseti-security/external_project_access_scanner_cron_runner.lock $FORSETI_HOME/install/gcp/scripts/run_external_project_access_scanner.sh || echo '[forseti-security] Warning: New external project access scanner cron job will not be started, because previous external project access scanner job is still running.') 2>&1 | logger") | crontab -u $USER -
    echo "Added the run_external_project_access_scanner.sh to crontab under user $USER"
fi

echo "Execution of startup script finished"
""".format(
    # Cloud SQL properties
    cloudsql_arch = context.properties['cloudsqlproxy-os-arch'],

    # Install Forseti.
    download_forseti=DOWNLOAD_FORSETI,

    # Checkout Forseti version.
    checkout_forseti_version=FORSETI_VERSION,

    # Set ownership for Forseti conf and rules dirs
    forseti_home=FORSETI_HOME,

    # Download the Forseti conf and rules.
    scanner_bucket=SCANNER_BUCKET,
    forseti_server_conf=FORSETI_SERVER_CONF,

    # Env variables for Explain
    export_initialize_vars=EXPORT_INITIALIZE_VARS,

    # Env variables for Forseti
    export_forseti_vars=EXPORT_FORSETI_VARS,

    # Forseti run frequency
    run_frequency=RUN_FREQUENCY,
    external_project_access_scanner_cron_schedule=EXTERNAL_PROJECT_ACCESS_SCANNER_CRON_SCHEDULE,
    external_project_access_scanner_cron_enabled=EXTERNAL_PROJECT_ACCESS_SCANNER_CRON_ENABLED,
)
                }]
            }
        }
    })
    return {'resources': resources}
<|endoftext|>


<commit_before>import os,sys,string
import math
import optparse
import apollo
import time
import glob,hashlib
import datetime
import random
from multiprocessing import Process,Queue

def parseFips(state,county,tract):
    fipsString = ""
    if len(str(state)) == 1:
        fipsString = "0" + str(state)
    else:
        fipsString = str(state)

    if len(str(county)) == 1:
        fipsString += "00" + str(county)
    elif len(str(county)) == 2:
        fipsString += "0" + str(county)
    else:
        fipsString += str(county)

    fipsString += str(tract)
    return fipsString

def error_exit(message):
	#statusUpdate("LOGERROR",message)
	sys.stderr.write(message)
	sys.exit(1)

if __name__ == '__main__':
    parser = optparse.OptionParser(usage="""
    %prog [--help]
    """)

    parser.add_option("-d","--debug",action="store_true")
    parser.add_option("-i","--runId",type="string",
		      help="The Apollo RunId for this call")

    opts,args=parser.parse_args()

    try:
        apolloDB = apollo.ApolloDB(host_="warhol-fred.psc.edu",dbname_="test")
        apolloDB.connect()

    except Exception as e:
        error_exit(str(e))

    tractList = []
    simulationRunId = opts.runId

    ### Parse the Tracts that you need to know
    try:
        with open("Tracts0","rb") as f:
            for line in f.readlines()[1:]:
                lineSplit = line.split(",")
                tractList.append(parseFips(lineSplit[1],lineSplit[2],lineSplit[3]))
    except Exception as e:
        error_exit(str(e))

    try:
        #### Gather op the list of Log files
        logFiles = glob.glob('Log*')

        ## First find the number of days that this job was run for
        numberOfDays = 0
        for logFile in logFiles:
            with open(logFile,"rb") as f:
            ## get the last line
                lastLineSplit = f.readlines()[-1].split(",")
                if int(lastLineSplit[0]) > numberOfDays: numberOfDays = int(lastLineSplit[0])

        numberOfDays += 1
        newInfectsPerDayPerTract = [[0.0 for x in range(numberOfDays)] for x in range(len(tractList))]

        for logFile in logFiles:

            with open(logFile,'rb') as f:
                ### skip Header
                isHead = False
                for line in f:
                    if not isHead:
                        isHead = True
                    else:
                        lineSplit = line.split(",")
                        newInfectsVal = sum(float(x) for x in lineSplit[12:16])
                        #print lineSplit[1] +" " + lineSplit[0]
                        newInfectsPerDayPerTract[int(lineSplit[1])][int(lineSplit[0])] += newInfectsVal

        ### Now make them an average
        for i in range(0,len(newInfectsPerDayPerTract)):
            for j in range(0,len(newInfectsPerDayPerTract[i])):
                newInfectsPerDayPerTract[i][j] /= float(len(logFiles))

        ### Create the output string
        outStrList = []
        for i in range(0,len(newInfectsPerDayPerTract)):
            for j in range(0,len(newInfectsPerDayPerTract[i])):
                outStrList.append('US%s %d %d:1\n'%(tractList[i],int(math.ceil(newInfectsPerDayPerTract[i][j])),j))

        outStr = ''.join(outStrList)
        m = hashlib.md5()
        m.update(outStr)
        m5hash = m.hexdigest()
	#m5hash = m5hash[:-2] + "1c"
        bufferSize=len(outStr)
        count = 0
        runDataContentId = -1
        SQLString = 'INSERT INTO run_data_content (text_content, md5_hash_of_content) values ("%s","%s")'%(outStr,m5hash)
        apolloDB.query(SQLString)
        runDataContentId = apolloDB.insertID()
        softIdTS =    apolloDB.getSoftwareIdentificationId(name_="Time Series Visualizer",version_="1.0")
        softIdGAIA =  apolloDB.getSoftwareIdentificationId(name_="GAIA",version_="1.0")
        softIdFLUTE = apolloDB.getSoftwareIdentificationId(name_="FluTE",version_="1.15")
        runDataDescriptionIdTS   = apolloDB.getRunDataDescriptionId(label_="newly_exposed.txt",
                                                                    source_software_=softIdFLUTE,
								    destination_software_=softIdTS)
        runDataDescriptionIdGAIA = apolloDB.getRunDataDescriptionId(label_="newly_exposed.txt",
                                                                    source_software_=softIdFLUTE,
                                                                    destination_software_=softIdGAIA)
        SQLString = 'INSERT INTO run_data(run_Id,description_id,content_id) values ("%s","%d","%d")'%(simulationRunId,
                                                                                runDataDescriptionIdTS,
                                                                                runDataContentId)
        apolloDB.query(SQLString)

        SQLString = 'INSERT INTO run_data(run_Id,description_id,content_id) values ("%s","%d","%d")'%(simulationRunId,
                                                                                runDataDescriptionIdGAIA,
                                                                                runDataContentId)
        apolloDB.query(SQLString)

        apolloDB.close()

        ### Insert data into database

        #with open('test.txt','wb') as f:
        #    f.write("%s"%(''.join(outStr)))

    except Exception as e:
        error_exit(str(e))

<commit_msg>Added the capability to reuse data_content that already exists upon checking the m5 hash.<commit_after>import os,sys,string
import math
import optparse
import apollo
import time
import glob,hashlib
import datetime
import random
from multiprocessing import Process,Queue

def parseFips(state,county,tract):
    fipsString = ""
    if len(str(state)) == 1:
        fipsString = "0" + str(state)
    else:
        fipsString = str(state)

    if len(str(county)) == 1:
        fipsString += "00" + str(county)
    elif len(str(county)) == 2:
        fipsString += "0" + str(county)
    else:
        fipsString += str(county)

    fipsString += str(tract)
    return fipsString

def error_exit(message):
	#statusUpdate("LOGERROR",message)
	sys.stderr.write(message)
	sys.exit(1)

if __name__ == '__main__':
    parser = optparse.OptionParser(usage="""
    %prog [--help]
    """)

    parser.add_option("-d","--debug",action="store_true")
    parser.add_option("-i","--runId",type="string",
		      help="The Apollo RunId for this call")

    opts,args=parser.parse_args()

    try:
        apolloDB = apollo.ApolloDB(host_="warhol-fred.psc.edu",dbname_="test")
        apolloDB.connect()

    except Exception as e:
        error_exit(str(e))

    tractList = []
    simulationRunId = opts.runId

    ### Parse the Tracts that you need to know
    try:
        with open("Tracts0","rb") as f:
            for line in f.readlines()[1:]:
                lineSplit = line.split(",")
                tractList.append(parseFips(lineSplit[1],lineSplit[2],lineSplit[3]))
    except Exception as e:
        error_exit(str(e))

    try:
        #### Gather op the list of Log files
        logFiles = glob.glob('Log*')

        ## First find the number of days that this job was run for
        numberOfDays = 0
        for logFile in logFiles:
            with open(logFile,"rb") as f:
            ## get the last line
                lastLineSplit = f.readlines()[-1].split(",")
                if int(lastLineSplit[0]) > numberOfDays: numberOfDays = int(lastLineSplit[0])

        numberOfDays += 1
        newInfectsPerDayPerTract = [[0.0 for x in range(numberOfDays)] for x in range(len(tractList))]

        for logFile in logFiles:

            with open(logFile,'rb') as f:
                ### skip Header
                isHead = False
                for line in f:
                    if not isHead:
                        isHead = True
                    else:
                        lineSplit = line.split(",")
                        newInfectsVal = sum(float(x) for x in lineSplit[12:16])
                        #print lineSplit[1] +" " + lineSplit[0]
                        newInfectsPerDayPerTract[int(lineSplit[1])][int(lineSplit[0])] += newInfectsVal

        ### Now make them an average
        for i in range(0,len(newInfectsPerDayPerTract)):
            for j in range(0,len(newInfectsPerDayPerTract[i])):
                newInfectsPerDayPerTract[i][j] /= float(len(logFiles))

        ### Create the output string
        outStrList = []
        for i in range(0,len(newInfectsPerDayPerTract)):
            for j in range(0,len(newInfectsPerDayPerTract[i])):
                outStrList.append('US%s %d %d:1\n'%(tractList[i],int(math.ceil(newInfectsPerDayPerTract[i][j])),j))

        outStr = ''.join(outStrList)
	with open('out.log','wb') as f:
	    	f.write(outStr)
        m = hashlib.md5()
        m.update(outStr)
        m5hash = m.hexdigest()
        bufferSize=len(outStr)
        count = 0
        runDataContentId = -1
        hashvar = apolloDB.checkMD5HashExistence(m5hash)
        if hashvar > -1:
            runDataContentId = hashvar
        else:
        	SQLString = 'INSERT INTO run_data_content (text_content, md5_hash_of_content) values ("%s","%s")'%(outStr,m5hash)
        	apolloDB.query(SQLString)
        	runDataContentId = apolloDB.insertID()
        softIdTS =    apolloDB.getSoftwareIdentificationId(name_="Time Series Visualizer",version_="1.0")
        softIdGAIA =  apolloDB.getSoftwareIdentificationId(name_="GAIA",version_="1.0")
        softIdFLUTE = apolloDB.getSoftwareIdentificationId(name_="FluTE",version_="1.15")
        runDataDescriptionIdTS   = apolloDB.getRunDataDescriptionId(label_="newly_exposed.txt",
                                                                    source_software_=softIdFLUTE,
								    destination_software_=softIdTS)
        runDataDescriptionIdGAIA = apolloDB.getRunDataDescriptionId(label_="newly_exposed.txt",
                                                                    source_software_=softIdFLUTE,
                                                                    destination_software_=softIdGAIA)
        SQLString = 'INSERT INTO run_data(run_Id,description_id,content_id) values ("%s","%d","%d")'%(simulationRunId,
                                                                                runDataDescriptionIdTS,
                                                                                runDataContentId)
        apolloDB.query(SQLString)

        SQLString = 'INSERT INTO run_data(run_Id,description_id,content_id) values ("%s","%d","%d")'%(simulationRunId,
                                                                                runDataDescriptionIdGAIA,
                                                                                runDataContentId)
        apolloDB.query(SQLString)

        apolloDB.close()

    except Exception as e:
        error_exit(str(e))

<|endoftext|>


<commit_before>from OpenGL.GL import *

from vistas.core.color import RGBColor
from vistas.core.graphics.matrix import ViewMatrix
from vistas.core.graphics.scene import Scene
from vistas.core.graphics.vector import Vector

class Camera:
    def __init__(self, scene=None, color=RGBColor(0, 0, 0)):
        if scene is None:
            scene = Scene()

        self.scene = scene
        self.color = color
        self.matrix = ViewMatrix()
        self.saved_matrix_state = None
        self.wireframe = False
        self.selection_view = False
        self.proj_matrix = ViewMatrix()

        self.set_up_vector(Vector(0, 1, 0))
        self.set_position(Vector(0, 0, 0))
        self.set_point_of_interest(Vector(0, 0, -10))

        # Todo: VI_CameraSyncObservable::GetInstance()->AddObserver(this);

    def __del__(self):
        pass  # Todo

    def get_position(self):
        relative_pos = Vector(self.matrix[3,0], self.matrix[3,1], self.matrix[3,2])
        relative_pos *= -1

        actual_pos = Vector(0, 0, 0)
        for i, attr in enumerate(['x','y','z']):
            setattr(
                actual_pos, attr,
                self.matrix[i,0]*relative_pos.x + self.matrix[i,1]*relative_pos.y + self.matrix[i,2]*relative_pos.z
            )

        return actual_pos

    def get_direction(self):
        return Vector(self.matrix[0,2]*-1, self.matrix[1,2]*-1, self.matrix[2,2]*-1)

    def set_point_of_interest(self, poi):
        pos = self.get_position()
        forward = poi - pos
        forward.normalize()

        right = forward.cross(self.get_up_vector())
        right.normalize()

        up = right.cross(forward)
        up.normalize()

        self.matrix[0,0] = right.x
        self.matrix[1,0] = right.y
        self.matrix[2,0] = right.z

        self.matrix[0,1] = up.x
        self.matrix[1,1] = up.y
        self.matrix[2,1] = up.z

        self.matrix[0,2] = forward.x * -1
        self.matrix[1,2] = forward.y * -1
        self.matrix[2,2] = forward.z * -1

        self.set_position(pos)

    def set_position(self, position):
        relative_pos = self.matrix * position * -1
        self.matrix[3,0] = relative_pos.x
        self.matrix[3,1] = relative_pos.y
        self.matrix[3,2] = relative_pos.z

    def get_up_vector(self):
        return Vector(self.matrix[0,1], self.matrix[1,1], self.matrix[2,1])

    def set_up_vector(self, up):
        unit_up = up.normalized
        pos = self.get_position()

        self.matrix[0,1] = unit_up.x
        self.matrix[1,1] = unit_up.y
        self.matrix[2,1] = unit_up.z

        right = self.get_direction().cross(unit_up)
        right.normalize()
        self.matrix[0,0] = right.x
        self.matrix[1,0] = right.y
        self.matrix[2,0] = right.z

        forward = unit_up.cross(right)
        forward.normalize()
        self.matrix[0,2] = forward.x * -1
        self.matrix[1,2] = forward.y * -1
        self.matrix[2,2] = forward.z * -1

        self.set_position(pos)

    def move_relative(self, movement):
        self.matrix *= ViewMatrix.translate(*movement.v[:3])

    def rotate_relative(self, rotation):
        pos = self.get_position()
        rotate_matrix = (
            ViewMatrix.rotate_x(rotation.x) * ViewMatrix.rotate_y(rotation.y) * ViewMatrix.rotate_z(rotation.z)
        )

        self.matrix = rotate_matrix * self.matrix
        self.set_position(pos)

    def distance_to_point(self, point):
        return abs((point - self.get_position()).length)

    def distance_to_object(self, obj):
        return abs((obj.bounds.center - self.get_position()).length)

    def render(self, width, height):
        if self.selection_view:
            self.reset(width, height, RGBColor(1, 1, 1, 1))
        else:
            self.reset(width, height, self.color)

        # Render scene
        if self.wireframe:
            glPolygonMode(GL_FRONT_AND_BACK, GL_LINE)
        else:
            glPolygonMode(GL_FRONT_AND_BACK, GL_FILL)

        if self.selection_view:
            self.scene.select_object(0, 0)
        else:
            self.scene.render(self)

    def render_to_bitmap(self, width, height):
        pass  # Todo

    def select_object(self, width, height, x, y):
        pass  # Todo

    def update(self, observable):
        pass # Todo

    def reset(self, width, height, color):
        scene_box = self.scene.bounding_box

        glEnable(GL_DEPTH_TEST)
        glDepthFunc(GL_LESS)
        glClearDepth(1.0)
        glClear(GL_DEPTH_BUFFER_BIT)
        glClearColor(color.r, color.g, color.b, color.a)
        glClear(GL_COLOR_BUFFER_BIT)

        glViewport(0, 0, width, height)

        #c = (self.matrix * Vector(*scene_box.center.v[:3], 1))     # Todo - fix z_near/z_far calculations or remove?
        #z_near = max(1, c.z - scene_box.diameter / 2)
        #z_far = (c.z + scene_box.diameter) * 2
        #self.proj_matrix = ViewMatrix.perspective(80.0, width / height, z_near, z_far)

        self.proj_matrix = ViewMatrix.perspective(80.0, width / height, 1, 10000)
<commit_msg>Use max float for z far value<commit_after>import sys
from OpenGL.GL import *

from vistas.core.color import RGBColor
from vistas.core.graphics.matrix import ViewMatrix
from vistas.core.graphics.scene import Scene
from vistas.core.graphics.vector import Vector

class Camera:
    def __init__(self, scene=None, color=RGBColor(0, 0, 0)):
        if scene is None:
            scene = Scene()

        self.scene = scene
        self.color = color
        self.matrix = ViewMatrix()
        self.saved_matrix_state = None
        self.wireframe = False
        self.selection_view = False
        self.proj_matrix = ViewMatrix()

        self.set_up_vector(Vector(0, 1, 0))
        self.set_position(Vector(0, 0, 0))
        self.set_point_of_interest(Vector(0, 0, -10))

        # Todo: VI_CameraSyncObservable::GetInstance()->AddObserver(this);

    def __del__(self):
        pass  # Todo

    def get_position(self):
        relative_pos = Vector(self.matrix[3,0], self.matrix[3,1], self.matrix[3,2])
        relative_pos *= -1

        actual_pos = Vector(0, 0, 0)
        for i, attr in enumerate(['x','y','z']):
            setattr(
                actual_pos, attr,
                self.matrix[i,0]*relative_pos.x + self.matrix[i,1]*relative_pos.y + self.matrix[i,2]*relative_pos.z
            )

        return actual_pos

    def get_direction(self):
        return Vector(self.matrix[0,2]*-1, self.matrix[1,2]*-1, self.matrix[2,2]*-1)

    def set_point_of_interest(self, poi):
        pos = self.get_position()
        forward = poi - pos
        forward.normalize()

        right = forward.cross(self.get_up_vector())
        right.normalize()

        up = right.cross(forward)
        up.normalize()

        self.matrix[0,0] = right.x
        self.matrix[1,0] = right.y
        self.matrix[2,0] = right.z

        self.matrix[0,1] = up.x
        self.matrix[1,1] = up.y
        self.matrix[2,1] = up.z

        self.matrix[0,2] = forward.x * -1
        self.matrix[1,2] = forward.y * -1
        self.matrix[2,2] = forward.z * -1

        self.set_position(pos)

    def set_position(self, position):
        relative_pos = self.matrix * position * -1
        self.matrix[3,0] = relative_pos.x
        self.matrix[3,1] = relative_pos.y
        self.matrix[3,2] = relative_pos.z

    def get_up_vector(self):
        return Vector(self.matrix[0,1], self.matrix[1,1], self.matrix[2,1])

    def set_up_vector(self, up):
        unit_up = up.normalized
        pos = self.get_position()

        self.matrix[0,1] = unit_up.x
        self.matrix[1,1] = unit_up.y
        self.matrix[2,1] = unit_up.z

        right = self.get_direction().cross(unit_up)
        right.normalize()
        self.matrix[0,0] = right.x
        self.matrix[1,0] = right.y
        self.matrix[2,0] = right.z

        forward = unit_up.cross(right)
        forward.normalize()
        self.matrix[0,2] = forward.x * -1
        self.matrix[1,2] = forward.y * -1
        self.matrix[2,2] = forward.z * -1

        self.set_position(pos)

    def move_relative(self, movement):
        self.matrix *= ViewMatrix.translate(*movement.v[:3])

    def rotate_relative(self, rotation):
        pos = self.get_position()
        rotate_matrix = (
            ViewMatrix.rotate_x(rotation.x) * ViewMatrix.rotate_y(rotation.y) * ViewMatrix.rotate_z(rotation.z)
        )

        self.matrix = rotate_matrix * self.matrix
        self.set_position(pos)

    def distance_to_point(self, point):
        return abs((point - self.get_position()).length)

    def distance_to_object(self, obj):
        return abs((obj.bounds.center - self.get_position()).length)

    def render(self, width, height):
        if self.selection_view:
            self.reset(width, height, RGBColor(1, 1, 1, 1))
        else:
            self.reset(width, height, self.color)

        # Render scene
        if self.wireframe:
            glPolygonMode(GL_FRONT_AND_BACK, GL_LINE)
        else:
            glPolygonMode(GL_FRONT_AND_BACK, GL_FILL)

        if self.selection_view:
            self.scene.select_object(0, 0)
        else:
            self.scene.render(self)

    def render_to_bitmap(self, width, height):
        pass  # Todo

    def select_object(self, width, height, x, y):
        pass  # Todo

    def update(self, observable):
        pass # Todo

    def reset(self, width, height, color):
        scene_box = self.scene.bounding_box

        glEnable(GL_DEPTH_TEST)
        glDepthFunc(GL_LESS)
        glClearDepth(1.0)
        glClear(GL_DEPTH_BUFFER_BIT)
        glClearColor(color.r, color.g, color.b, color.a)
        glClear(GL_COLOR_BUFFER_BIT)

        glViewport(0, 0, width, height)

        #c = (self.matrix * Vector(*scene_box.center.v[:3], 1))     # Todo - fix z_near/z_far calculations or remove?
        #z_near = max(1, c.z - scene_box.diameter / 2)
        #z_far = (c.z + scene_box.diameter) * 2
        #self.proj_matrix = ViewMatrix.perspective(80.0, width / height, z_near, z_far)

        self.proj_matrix = ViewMatrix.perspective(80.0, width / height, 1, sys.float_info.max)
<|endoftext|>


<commit_before>"""
    :codeauthor: Mike Place <mp@saltstack.com>
"""

import copy
import os
import pathlib
import shutil
import tempfile
import textwrap

import salt.fileclient
import salt.fileserver.roots as roots
import salt.utils.files
import salt.utils.hashutils
import salt.utils.platform
import salt.utils.stringutils
from tests.support.mixins import (
    AdaptedConfigurationTestCaseMixin,
    LoaderModuleMockMixin,
)
from tests.support.mock import MagicMock, mock_open, patch
from tests.support.runtests import RUNTIME_VARS
from tests.support.unit import TestCase, skipIf

UNICODE_FILENAME = "питон.txt"
UNICODE_DIRNAME = UNICODE_ENVNAME = "соль"

class RootsTest(TestCase, AdaptedConfigurationTestCaseMixin, LoaderModuleMockMixin):
    def setup_loader_modules(self):
        config_overrides = {"file_roots": {"base": [str(self.tmp_state_tree)]}}
        self.opts = self.get_temp_config("master", **config_overrides)
        return {roots: {"__opts__": self.opts}}

    @classmethod
    def setUpClass(cls):
        cls.tmp_dir = pathlib.Path(tempfile.mkdtemp(dir=RUNTIME_VARS.TMP))
        cls.tmp_state_tree = pathlib.Path(tempfile.mkdtemp(dir=RUNTIME_VARS.TMP))
        full_path_to_file = os.path.join(RUNTIME_VARS.BASE_FILES, "testfile")
        shutil.copyfile(full_path_to_file, str(cls.tmp_dir / "testfile"))
        shutil.copyfile(full_path_to_file, str(cls.tmp_state_tree / "testfile"))
        shutil.copyfile(
            os.path.join(RUNTIME_VARS.BASE_FILES, UNICODE_FILENAME),
            str(cls.tmp_state_tree / UNICODE_FILENAME),
        )
        shutil.copytree(
            os.path.join(RUNTIME_VARS.BASE_FILES, UNICODE_DIRNAME),
            str(cls.tmp_state_tree / UNICODE_DIRNAME),
        )

    @classmethod
    def tearDownClass(cls):
        salt.utils.files.rm_rf(str(cls.tmp_dir))
        salt.utils.files.rm_rf(str(cls.tmp_state_tree))

    def tearDown(self):
        del self.opts

    def test_file_list(self):
        ret = roots.file_list({"saltenv": "base"})
        self.assertIn("testfile", ret)
        self.assertIn(UNICODE_FILENAME, ret)

    def test_find_file(self):
        ret = roots.find_file("testfile")
        self.assertEqual("testfile", ret["rel"])

        full_path_to_file = str(self.tmp_state_tree / "testfile")
        self.assertEqual(full_path_to_file, ret["path"])

    def test_serve_file(self):
        with patch.dict(roots.__opts__, {"file_buffer_size": 262144}):
            load = {
                "saltenv": "base",
                "path": str(self.tmp_dir / "testfile"),
                "loc": 0,
            }
            fnd = {"path": str(self.tmp_dir / "testfile"), "rel": "testfile"}
            ret = roots.serve_file(load, fnd)

            with salt.utils.files.fopen(
                os.path.join(RUNTIME_VARS.BASE_FILES, "testfile"), "rb"
            ) as fp_:
                data = fp_.read()

            self.assertDictEqual(ret, {"data": data, "dest": "testfile"})

    def test_envs(self):
        opts = {"file_roots": copy.copy(self.opts["file_roots"])}
        opts["file_roots"][UNICODE_ENVNAME] = opts["file_roots"]["base"]
        with patch.dict(roots.__opts__, opts):
            ret = roots.envs()
        self.assertIn("base", ret)
        self.assertIn(UNICODE_ENVNAME, ret)

    def test_file_hash(self):
        load = {
            "saltenv": "base",
            "path": str(self.tmp_dir / "testfile"),
        }
        fnd = {"path": str(self.tmp_dir / "testfile"), "rel": "testfile"}
        ret = roots.file_hash(load, fnd)

        # Hashes are different in Windows. May be how git translates line
        # endings
        with salt.utils.files.fopen(
            os.path.join(RUNTIME_VARS.BASE_FILES, "testfile"), "rb"
        ) as fp_:
            hsum = salt.utils.hashutils.sha256_digest(fp_.read())

        self.assertDictEqual(ret, {"hsum": hsum, "hash_type": "sha256"})

    def test_file_list_emptydirs(self):
        empty_dir = self.tmp_state_tree / "empty_dir"
        if not empty_dir.is_dir():
            empty_dir.mkdir()
        self.addCleanup(salt.utils.files.rm_rf, str(empty_dir))
        ret = roots.file_list_emptydirs({"saltenv": "base"})
        self.assertIn("empty_dir", ret)

    def test_file_list_with_slash(self):
        opts = {"file_roots": copy.copy(self.opts["file_roots"])}
        opts["file_roots"]["foo/bar"] = opts["file_roots"]["base"]
        load = {
            "saltenv": "foo/bar",
        }
        with patch.dict(roots.__opts__, opts):
            ret = roots.file_list(load)
        self.assertIn("testfile", ret)
        self.assertIn(UNICODE_FILENAME, ret)

    def test_dir_list(self):
        empty_dir = self.tmp_state_tree / "empty_dir"
        if not empty_dir.is_dir():
            empty_dir.mkdir()
        self.addCleanup(salt.utils.files.rm_rf, str(empty_dir))
        ret = roots.dir_list({"saltenv": "base"})
        self.assertIn("empty_dir", ret)
        self.assertIn(UNICODE_DIRNAME, ret)

    def test_symlink_list(self):
        source_sym = self.tmp_state_tree / "source_sym"
        source_sym.write_text("")
        dest_sym = self.tmp_state_tree / "dest_sym"
        dest_sym.symlink_to(str(source_sym))
        self.addCleanup(dest_sym.unlink)
        self.addCleanup(source_sym.unlink)
        ret = roots.symlink_list({"saltenv": "base"})
        self.assertDictEqual(ret, {"dest_sym": str(source_sym)})

    def test_dynamic_file_roots(self):
        dyn_root_dir = tempfile.mkdtemp(dir=RUNTIME_VARS.TMP)
        top_sls = os.path.join(dyn_root_dir, "top.sls")
        with salt.utils.files.fopen(top_sls, "w") as fp_:
            fp_.write("{{saltenv}}:\n  '*':\n    - dynamo\n")
        dynamo_sls = os.path.join(dyn_root_dir, "dynamo.sls")
        with salt.utils.files.fopen(dynamo_sls, "w") as fp_:
            fp_.write("foo:\n  test.nop\n")
        opts = {"file_roots": copy.copy(self.opts["file_roots"])}
        opts["file_roots"]["__env__"] = [dyn_root_dir]
        with patch.dict(roots.__opts__, opts):
            ret1 = roots.find_file("dynamo.sls", "dyn")
            ret2 = roots.file_list({"saltenv": "dyn"})
        self.assertEqual("dynamo.sls", ret1["rel"])
        self.assertIn("top.sls", ret2)
        self.assertIn("dynamo.sls", ret2)

    @skipIf(
        salt.utils.platform.is_windows(),
        "Windows does not support this master function",
    )
    def test_update_no_change(self):
        # process all changes that have happen
        # changes will always take place the first time during testing
        ret = roots.update()
        self.assertTrue(ret["changed"])

        # check if no changes took place
        ret = roots.update()
        self.assertFalse(ret["changed"])
        self.assertEqual(ret["files"]["changed"], [])
        self.assertEqual(ret["files"]["removed"], [])
        self.assertEqual(ret["files"]["added"], [])

    def test_update_mtime_map(self):
        """
        Test that files with colons in the filename are properly handled in the
        mtime_map, and that they are properly identified as having changed.
        """
        mtime_map_path = os.path.join(self.opts["cachedir"], "roots", "mtime_map")
        mtime_map_mock = mock_open(
            read_data={
                mtime_map_path: textwrap.dedent(
                    """\
                    /srv/salt/kleine_Datei.txt:1594263154.0469685
                    /srv/salt/große:Datei.txt:1594263160.9336357
                    """
                ),
            }
        )
        new_mtime_map = {
            "/srv/salt/kleine_Datei.txt": 1594263154.0469685,
            "/srv/salt/große:Datei.txt": 1594263261.0616212,
        }

        with patch(
            "salt.fileserver.reap_fileserver_cache_dir", MagicMock(return_value=True)
        ), patch(
            "salt.fileserver.generate_mtime_map", MagicMock(return_value=new_mtime_map)
        ), patch.dict(
            roots.__opts__, {"fileserver_events": False}
        ), patch(
            "salt.utils.files.fopen", mtime_map_mock
        ):
            ret = roots.update()

        # Confirm the expected return from the function
        assert ret == {
            "changed": True,
            "files": {
                "changed": ["/srv/salt/große:Datei.txt"],
                "removed": [],
                "added": [],
            },
            "backend": "roots",
        }, ret

        # Confirm that the new values were written to the mtime_map. Sort both
        # lists of lines to account for variances in dictionary iteration order
        # between Python releases.
        lines_written = sorted(mtime_map_mock.write_calls())
        expected = sorted(
            [
                salt.utils.stringutils.to_bytes(
                    "{key}:{val}\n".format(key=key, val=val)
                )
                for key, val in new_mtime_map.items()
            ]
        )
        assert lines_written == expected, lines_written

    def test_update_mtime_map_unicode_error(self):
        """
        Test that a malformed mtime_map (which causes an UnicodeDecodeError
        exception) is handled properly.
        """
        new_mtime_map = {
            "/srv/salt/große:Datei.txt": 1594263261.0616212,
        }
        with tempfile.TemporaryDirectory() as tmpdirname:
            mtime_map_path = os.path.join(tmpdirname, "roots", "mtime_map")
            os.makedirs(os.path.dirname(mtime_map_path))
            with salt.utils.files.fopen(mtime_map_path, "wb") as fp:
                fp.write(b"\x9c")

            with patch(
                "salt.fileserver.reap_fileserver_cache_dir",
                MagicMock(return_value=True)
            ), patch(
                "salt.fileserver.generate_mtime_map",
                MagicMock(return_value=new_mtime_map)
            ), patch.dict(
                roots.__opts__,
                {"fileserver_events": False, "cachedir": tmpdirname},
            ):
                ret = roots.update()

        assert ret == {
            "changed": True,
            "files": {
                "changed": [],
                "removed": [],
                "added": ["/srv/salt/große:Datei.txt"],
            },
            "backend": "roots",
        }, ret
<commit_msg>Update test_roots.py<commit_after>"""
    :codeauthor: Mike Place <mp@saltstack.com>
"""

import copy
import os
import pathlib
import shutil
import tempfile
import textwrap

import salt.fileclient
import salt.fileserver.roots as roots
import salt.utils.files
import salt.utils.hashutils
import salt.utils.platform
import salt.utils.stringutils
from tests.support.mixins import (
    AdaptedConfigurationTestCaseMixin,
    LoaderModuleMockMixin,
)
from tests.support.mock import MagicMock, mock_open, patch
from tests.support.runtests import RUNTIME_VARS
from tests.support.unit import TestCase, skipIf

UNICODE_FILENAME = "питон.txt"
UNICODE_DIRNAME = UNICODE_ENVNAME = "соль"

class RootsTest(TestCase, AdaptedConfigurationTestCaseMixin, LoaderModuleMockMixin):
    def setup_loader_modules(self):
        config_overrides = {"file_roots": {"base": [str(self.tmp_state_tree)]}}
        self.opts = self.get_temp_config("master", **config_overrides)
        return {roots: {"__opts__": self.opts}}

    @classmethod
    def setUpClass(cls):
        cls.tmp_dir = pathlib.Path(tempfile.mkdtemp(dir=RUNTIME_VARS.TMP))
        cls.tmp_state_tree = pathlib.Path(tempfile.mkdtemp(dir=RUNTIME_VARS.TMP))
        full_path_to_file = os.path.join(RUNTIME_VARS.BASE_FILES, "testfile")
        shutil.copyfile(full_path_to_file, str(cls.tmp_dir / "testfile"))
        shutil.copyfile(full_path_to_file, str(cls.tmp_state_tree / "testfile"))
        shutil.copyfile(
            os.path.join(RUNTIME_VARS.BASE_FILES, UNICODE_FILENAME),
            str(cls.tmp_state_tree / UNICODE_FILENAME),
        )
        shutil.copytree(
            os.path.join(RUNTIME_VARS.BASE_FILES, UNICODE_DIRNAME),
            str(cls.tmp_state_tree / UNICODE_DIRNAME),
        )

    @classmethod
    def tearDownClass(cls):
        salt.utils.files.rm_rf(str(cls.tmp_dir))
        salt.utils.files.rm_rf(str(cls.tmp_state_tree))

    def tearDown(self):
        del self.opts

    def test_file_list(self):
        ret = roots.file_list({"saltenv": "base"})
        self.assertIn("testfile", ret)
        self.assertIn(UNICODE_FILENAME, ret)

    def test_find_file(self):
        ret = roots.find_file("testfile")
        self.assertEqual("testfile", ret["rel"])

        full_path_to_file = str(self.tmp_state_tree / "testfile")
        self.assertEqual(full_path_to_file, ret["path"])

    def test_serve_file(self):
        with patch.dict(roots.__opts__, {"file_buffer_size": 262144}):
            load = {
                "saltenv": "base",
                "path": str(self.tmp_dir / "testfile"),
                "loc": 0,
            }
            fnd = {"path": str(self.tmp_dir / "testfile"), "rel": "testfile"}
            ret = roots.serve_file(load, fnd)

            with salt.utils.files.fopen(
                os.path.join(RUNTIME_VARS.BASE_FILES, "testfile"), "rb"
            ) as fp_:
                data = fp_.read()

            self.assertDictEqual(ret, {"data": data, "dest": "testfile"})

    def test_envs(self):
        opts = {"file_roots": copy.copy(self.opts["file_roots"])}
        opts["file_roots"][UNICODE_ENVNAME] = opts["file_roots"]["base"]
        with patch.dict(roots.__opts__, opts):
            ret = roots.envs()
        self.assertIn("base", ret)
        self.assertIn(UNICODE_ENVNAME, ret)

    def test_file_hash(self):
        load = {
            "saltenv": "base",
            "path": str(self.tmp_dir / "testfile"),
        }
        fnd = {"path": str(self.tmp_dir / "testfile"), "rel": "testfile"}
        ret = roots.file_hash(load, fnd)

        # Hashes are different in Windows. May be how git translates line
        # endings
        with salt.utils.files.fopen(
            os.path.join(RUNTIME_VARS.BASE_FILES, "testfile"), "rb"
        ) as fp_:
            hsum = salt.utils.hashutils.sha256_digest(fp_.read())

        self.assertDictEqual(ret, {"hsum": hsum, "hash_type": "sha256"})

    def test_file_list_emptydirs(self):
        empty_dir = self.tmp_state_tree / "empty_dir"
        if not empty_dir.is_dir():
            empty_dir.mkdir()
        self.addCleanup(salt.utils.files.rm_rf, str(empty_dir))
        ret = roots.file_list_emptydirs({"saltenv": "base"})
        self.assertIn("empty_dir", ret)

    def test_file_list_with_slash(self):
        opts = {"file_roots": copy.copy(self.opts["file_roots"])}
        opts["file_roots"]["foo/bar"] = opts["file_roots"]["base"]
        load = {
            "saltenv": "foo/bar",
        }
        with patch.dict(roots.__opts__, opts):
            ret = roots.file_list(load)
        self.assertIn("testfile", ret)
        self.assertIn(UNICODE_FILENAME, ret)

    def test_dir_list(self):
        empty_dir = self.tmp_state_tree / "empty_dir"
        if not empty_dir.is_dir():
            empty_dir.mkdir()
        self.addCleanup(salt.utils.files.rm_rf, str(empty_dir))
        ret = roots.dir_list({"saltenv": "base"})
        self.assertIn("empty_dir", ret)
        self.assertIn(UNICODE_DIRNAME, ret)

    def test_symlink_list(self):
        source_sym = self.tmp_state_tree / "source_sym"
        source_sym.write_text("")
        dest_sym = self.tmp_state_tree / "dest_sym"
        dest_sym.symlink_to(str(source_sym))
        self.addCleanup(dest_sym.unlink)
        self.addCleanup(source_sym.unlink)
        ret = roots.symlink_list({"saltenv": "base"})
        self.assertDictEqual(ret, {"dest_sym": str(source_sym)})

    def test_dynamic_file_roots(self):
        dyn_root_dir = tempfile.mkdtemp(dir=RUNTIME_VARS.TMP)
        top_sls = os.path.join(dyn_root_dir, "top.sls")
        with salt.utils.files.fopen(top_sls, "w") as fp_:
            fp_.write("{{saltenv}}:\n  '*':\n    - dynamo\n")
        dynamo_sls = os.path.join(dyn_root_dir, "dynamo.sls")
        with salt.utils.files.fopen(dynamo_sls, "w") as fp_:
            fp_.write("foo:\n  test.nop\n")
        opts = {"file_roots": copy.copy(self.opts["file_roots"])}
        opts["file_roots"]["__env__"] = [dyn_root_dir]
        with patch.dict(roots.__opts__, opts):
            ret1 = roots.find_file("dynamo.sls", "dyn")
            ret2 = roots.file_list({"saltenv": "dyn"})
        self.assertEqual("dynamo.sls", ret1["rel"])
        self.assertIn("top.sls", ret2)
        self.assertIn("dynamo.sls", ret2)

    @skipIf(
        salt.utils.platform.is_windows(),
        "Windows does not support this master function",
    )
    def test_update_no_change(self):
        # process all changes that have happen
        # changes will always take place the first time during testing
        ret = roots.update()
        self.assertTrue(ret["changed"])

        # check if no changes took place
        ret = roots.update()
        self.assertFalse(ret["changed"])
        self.assertEqual(ret["files"]["changed"], [])
        self.assertEqual(ret["files"]["removed"], [])
        self.assertEqual(ret["files"]["added"], [])

    def test_update_mtime_map(self):
        """
        Test that files with colons in the filename are properly handled in the
        mtime_map, and that they are properly identified as having changed.
        """
        mtime_map_path = os.path.join(self.opts["cachedir"], "roots", "mtime_map")
        mtime_map_mock = mock_open(
            read_data={
                mtime_map_path: textwrap.dedent(
                    """\
                    /srv/salt/kleine_Datei.txt:1594263154.0469685
                    /srv/salt/große:Datei.txt:1594263160.9336357
                    """
                ),
            }
        )
        new_mtime_map = {
            "/srv/salt/kleine_Datei.txt": 1594263154.0469685,
            "/srv/salt/große:Datei.txt": 1594263261.0616212,
        }

        with patch(
            "salt.fileserver.reap_fileserver_cache_dir", MagicMock(return_value=True)
        ), patch(
            "salt.fileserver.generate_mtime_map", MagicMock(return_value=new_mtime_map)
        ), patch.dict(
            roots.__opts__, {"fileserver_events": False}
        ), patch(
            "salt.utils.files.fopen", mtime_map_mock
        ):
            ret = roots.update()

        # Confirm the expected return from the function
        assert ret == {
            "changed": True,
            "files": {
                "changed": ["/srv/salt/große:Datei.txt"],
                "removed": [],
                "added": [],
            },
            "backend": "roots",
        }, ret

        # Confirm that the new values were written to the mtime_map. Sort both
        # lists of lines to account for variances in dictionary iteration order
        # between Python releases.
        lines_written = sorted(mtime_map_mock.write_calls())
        expected = sorted(
            [
                salt.utils.stringutils.to_bytes(
                    "{key}:{val}\n".format(key=key, val=val)
                )
                for key, val in new_mtime_map.items()
            ]
        )
        assert lines_written == expected, lines_written

    def test_update_mtime_map_unicode_error(self):
        """
        Test that a malformed mtime_map (which causes an UnicodeDecodeError
        exception) is handled properly.
        """
        new_mtime_map = {
            "/srv/salt/große:Datei.txt": 1594263261.0616212,
        }
        with tempfile.TemporaryDirectory() as tmpdirname:
            mtime_map_path = os.path.join(tmpdirname, "roots", "mtime_map")
            os.makedirs(os.path.dirname(mtime_map_path))
            with salt.utils.files.fopen(mtime_map_path, "wb") as fp:
                fp.write(b"\x9c")

            with patch(
                "salt.fileserver.reap_fileserver_cache_dir",
                MagicMock(return_value=True)
            ), patch(
                "salt.fileserver.generate_mtime_map",
                MagicMock(return_value=new_mtime_map)
            ), patch.dict(
                roots.__opts__,
                {"fileserver_events": False, "cachedir": tmpdirname},
            ):
                ret = roots.update()

        assert ret == {
            "changed": True,
            "files": {
                "changed": [],
                "removed": [],
                "added": ["/srv/salt/große:Datei.txt"],
            },
            "backend": "roots",
        }, ret
<|endoftext|>