Customizing pyATS – AeTest – Part 2

In part 1, we created a yaml file with some custom data. We mapped that into a python dictionary within a module called custombits which we created. We then imported that into our AeTest job file so we could pick and choose which jobs we ran and the ordering of those jobs. Before reading on make sure you have reviewed part 1 or this could be difficult to follow.

That was all pretty cool but we might want to do the same for the actual scripts themselves i.e. we don’t want to hard code ping vrf default 192.168.1.1 into our script. It would be more useful if we could provide a list of customizable ping destinations, VRFs and also the count. We might want more than the default number of ping replies.

As a starting point let’s look back at an extract from the customize.yaml file created in Part 1.

Ping_Dest:

 dest1 :
    index: 1 
    vrf: default  
    destination: 192.168.1.254
    count: 10

 dest2 :
    index: 2 
    vrf: default  
    destination: 192.168.1.212 
    count: 5

 dest3 :
    index: 3 
    vrf: default  
    destination: 192.168.1.214
    count: 50

All of this data will be mapped into our AeTest Script. Let’s just put in the full AeTest PING script below and examine what’s going on.

"""
ping_test.py

"""

import logging
import os
from pyats import aetest
from genie.testbed import load
from genie import parsergen
import re
from unicon.core.errors import TimeoutError, StateMachineError, ConnectionError
from custombits import CustomBits


# create a logger for this module
logger = logging.getLogger(__name__)


class CommonSetup(aetest.CommonSetup):
    @aetest.subsection
    def load_testbed(self, testbed):
        # Convert pyATS testbed to Genie Testbed
        logger.info(
            "Converting pyATS testbed to Genie Testbed to support pyATS Library features"
        )
        testbed = load(testbed)
        self.parent.parameters.update(testbed=testbed)

    @aetest.subsection
    def connect(self, testbed):
        """
        establishes connection to all your testbed devices.
        """
        # make sure testbed is provided
        assert testbed, "Testbed is not provided!"

        try:
            testbed.connect()
        except (TimeoutError, StateMachineError, ConnectionError):
            logger.error("Unable to connect to all devices")


class connectivity_status_checks(aetest.Testcase):


    @aetest.setup
    def setup(self, testbed):
        """Learn and save the interface details from the testbed devices."""
        self.execute_ping = {}
        self.output= {}
        self.cust_ping= CustomBits().ping_dest()
        for device_name, device in testbed.devices.items():
            
            if device.os in ("iosxr"):
                logger.info(f"{device_name} connected status: {device.connected}")
                logger.info(f"Running the PING command for {device_name}")
                for k,v in self.cust_ping.items():
                    vrf_name = v.get('vrf')
                    dest  = v.get('destination')
                    ping_cnt  = v.get('count')  
                    self.execute_ping[device_name, k] = device.execute(f"ping vrf {vrf_name} {dest} count {ping_cnt}")        

                



    @aetest.test
    def test(self, steps):
        for device_name, device in self.execute_ping.items():
                        with steps.start(
                            f"Running the PING command on {device_name}", continue_=True
                        ) as device_step:
            
                            output = self.execute_ping[device_name]
                            output = output.replace('\r\n','\n')
                            match = re.search(r'Success rate is (?P<rate>\d+) percent', output)
                            success_rate = match.group('rate')
                            logger.info(f' The success rate of the PING was {success_rate}')
                            if success_rate == 100:
                                
                                device_step.passed(f' No packets dropped, connectivity is good')
                                
                                
                            else:
                                device_step.failed(f'We had a problem with some of these tests')
        
            

class CommonCleanup(aetest.CommonCleanup):
    """CommonCleanup Section

    < common cleanup docstring >

    """

    # uncomment to add new subsections
    # @aetest.subsection
    # def subsection_cleanup_one(self):
    #     pass


if __name__ == "__main__":
    # for stand-alone execution
    import argparse
    from pyats import topology

    # from genie.conf import Genie

    parser = argparse.ArgumentParser(description="standalone parser")
    parser.add_argument(
        "--testbed",
        dest="testbed",
        help="testbed YAML file",
        type=topology.loader.load,
        # type=Genie.init,
        default=None,
    )

    # do the parsing
    args = parser.parse_known_args()[0]

    aetest.main(testbed=args.testbed)

As before we need to import our module:

from custombits import CustomBits

Starting in the setup portion of the script, we need to call our module and we map that into a variable called self.cust_ping. Note the use of self since AeTest uses classes. Also use of self ensures that the variable is callable within other sections such as the Test section will follows the setup section.

self.cust_ping= CustomBits().ping_dest()

We need a for loop to iterate over all the PING destinations and execute the commands. Note I am using IOS-XR here again in my lab.

                for k,v in self.cust_ping.items():
                    vrf_name = v.get('vrf')
                    dest  = v.get('destination')
                    ping_cnt  = v.get('count')  
                    self.execute_ping[device_name, k] = device.execute(f"ping vrf {vrf_name} {dest} count {ping_cnt}")    

We pull out all the different fields using the in-built “get“. We need to do a couple of things on the command execution line. First, all the output collected is mapped into a dictionary but we need two items here. We need the device name and key with the Dest1/Dest2 or whatever name we called it within the yaml file. Hence we have [device_name, k] and not just [device_name]. The “k” shorthand for “key” is a string which is pulled out from the For Loop into our overall dictionary. Using the for loop we will start to unpack the keys and values then map those into our command which will be executed on the device.

We use a f-string here for the placeholders for the variables to map into the command i.e. {dest}, {ping_cnt}, {vrf_name}.

Hopefully, this makes sense in terms of us pulling the data from the yaml file via the custom module and then placing this within our script. Now we could add more PING destinations, different VRFs and so on.

The actual test section, I used a simple REGEX instead of a Parser for this one. PING output is relatively simple on Cisco IOS-XR. We make use of “Step” in the test section so that we get a result per PING instead of one big pass/fail test.

In summary, this is probably not the only way to do this but this is the method that works for me. I really like to work with YAML instead of constantly tweaking scripts. This is especially useful when I am doing testing. Otherwise you end of with hundreds of variations of scripts instead of a single source of truth (i.e. a yaml file with custom data)

Published by gwoodwa1

IP Network Design and coding hobbyist

Leave a comment

Design a site like this with WordPress.com
Get started