Automatize sua rede com pyATS (DEMO) Pt2

# #####################################################################
# ####                       arquivo pyATS.py                       ###
# #####################################################################

#!/bin/env python

# Módulo para capturar log durante a execução do script
import logging
import json

# Irá construir a tabela onde será alocado dados da feature BGP
from tabulate import tabulate

# Necessário para o parâmetro aetest
from ats import aetest
from ats.log.utils import banner

# Importando módulos Genie
from genie.conf import Genie
from genie.abstract import Lookup

# Importando a library Genie
from genie.libs import ops 

log = logging.getLogger(__name__)

class common_setup(aetest.CommonSetup):
    """ Common Setup section """

    # CommonSetup pode haver diversas funções alocadas dentro da classe.
    
    # Definindo conexão com cada device descritos na mesa de teste que seria o arquivo testbed.yaml.
    # No caso, temos apenas um device descrito na mesa de teste.
    @aetest.subsection
    def connect(self, testbed):
        genie_testbed = Genie.init(testbed)
        self.parent.parameters['testbed'] = genie_testbed
        device_list = []
        for device in genie_testbed.devices.values():
            log.info(banner(
                "Connect to device '{d}'".format(d=device.name)))
            try:
                device.connect()
            except Exception as e:
                self.failed("Failed to establish connection to '{}'".format(
                    device.name))

            device_list.append(device)

        # Passando a lista de dispositivos para a variável "dev"
        self.parent.parameters.update(dev=device_list)

class BGP_Neighbors_Established(aetest.Testcase):
    """ Abrindo sessão de test case """

    # First test section
    @ aetest.test
    def learn_bgp(self):
        """ Montando estrutura """
        
        # all_bgp_sessions se refere ao módulo que faz parte da library pyATS
        self.all_bgp_sessions = {}
        for dev in self.parent.parameters['dev']:
            log.info(banner("Coletando informações da feature BGP {}".format(
                dev.name)))
            abstract = Lookup.from_device(dev)
            bgp = abstract.ops.bgp.bgp.Bgp(dev)
            bgp.learn()
            self.all_bgp_sessions[dev.name] = bgp.info

    @ aetest.test
    def check_bgp(self):
        """ Manipulando dados capturados da feature """

        failed_dict = {}
        mega_tabular = []
        for device, bgp in self.all_bgp_sessions.items():
            # Os valores abaixo definidos como string se refere aos parâmetros que estão alocados dentro dos comandos do BGP
            # A ideia de aprender a feature BGP é fazer com que todos os parâmetros das saída dos comandos 
            relacionados ao BGP sejam armazenados dentro da variável dev.name
            default = bgp['instance']['default']['vrf']['default']
            neighbors = default['neighbor']
            for nbr, props in neighbors.items():
                state = props.get('session_state')
                if state:
                    tr = []
                    tr.append(device)
                    tr.append(nbr)
                    tr.append(state)
                    if state == 'established' or state == 'Established':
                        tr.append('Passed')
                    else:
                        failed_dict[device] = {}
                        failed_dict[device][nbr] = props
                        tr.append('Failed')

                mega_tabular.append(tr)

        log.info(tabulate(mega_tabular,
                          headers=['Device', 'Peer',
                                   'State', 'Pass/Fail'],
                          tablefmt='orgtbl'))

        if failed_dict:
            log.error(json.dumps(failed_dict, indent=3))
            self.failed("Testbed has BGP Neighbors that are not established")

        else:
            self.passed("All BGP Neighbors are established")

class common_cleanup(aetest.CommonCleanup):
    """ Common Cleanup """
    @aetest.subsection
    def clean_everything(self):
        """ Common Cleanup Subsection """
        log.info("Aetest Common Cleanup ")

if __name__ == '__main__':  
    aetest.main()

# #####################################################################
# ####                       arquivo job  .py                       ###
# #####################################################################

import os
from ats.easypy import run

# A função abaixo irá rodar o script pyTAS.py
def main():
    # Find the location of the script in relation to the job file
    pwd = os.path.dirname(__file__)
    bgp_tests = os.path.join(pwd, 'pyTAS.py')
    # Execute the testscript
    # run(testscript=testscript)
    run(testscript=bgp_tests)