Merge "[ONOS-7636] Add function to verify host locations"
diff --git a/TestON/drivers/common/cli/onosclidriver.py b/TestON/drivers/common/cli/onosclidriver.py
index 117a593..7af7e4f 100755
--- a/TestON/drivers/common/cli/onosclidriver.py
+++ b/TestON/drivers/common/cli/onosclidriver.py
@@ -1151,6 +1151,58 @@
             main.log.exception( self.name + ": Uncaught exception!" )
             main.cleanAndExit()
 
+    def verifyHostLocation( self, hostIp, location ):
+        """
+        Description:
+            Verify the host given is discovered in all locations expected
+        Required:
+            hostIp: IP address of the host
+            location: expected location(s) of the given host. ex. "of:0000000000000005/8"
+                      Could be a string or list
+        Returns:
+            main.TRUE if host is discovered on all locations provided
+            main.FALSE otherwise
+        """
+        import json
+        locations = [ location ] if isinstance( location, str ) else location
+        assert isinstance( locations, list ), "Wrong type of location: {}".format( type( location ) )
+        try:
+            hosts = self.hosts()
+            hosts = json.loads( hosts )
+            targetHost = None
+            for host in hosts:
+                if hostIp in host[ "ipAddresses" ]:
+                    targetHost = host
+            assert host, "Not able to find host with IP {}".format( hostIp )
+            result = main.TRUE
+            locationsDiscovered = [ loc[ "elementId" ] + "/" + loc[ "port" ] for loc in targetHost[ "locations" ] ]
+            for loc in locations:
+                discovered = False
+                for locDiscovered in locationsDiscovered:
+                    if loc in locDiscovered:
+                        discovered = True
+                        main.log.debug( "Host {} discovered with location {}".format( hostIp, loc ) )
+                        break
+                if discovered:
+                    locationsDiscovered.remove( locDiscovered )
+                else:
+                    main.log.warn( "Host {} not discovered with location {}".format( hostIp, loc ) )
+                    result = main.FALSE
+            if locationsDiscovered:
+                main.log.warn( "Host {} is also discovered with location {}".format( hostIp, locationsDiscovered ) )
+                result = main.FALSE
+            return result
+        except KeyError:
+            main.log.exception( self.name + ": host data not as expected: " + hosts )
+            return None
+        except pexpect.EOF:
+            main.log.error( self.name + ": EOF exception found" )
+            main.log.error( self.name + ":     " + self.handle.before )
+            main.cleanAndExit()
+        except Exception:
+            main.log.exception( self.name + ": Uncaught exception" )
+            return None
+
     def verifyHostIp( self, hostList=[], prefix="" ):
         """
         Description:
diff --git a/TestON/tests/USECASE/SegmentRouting/dependencies/Testcaselib.py b/TestON/tests/USECASE/SegmentRouting/dependencies/Testcaselib.py
index 5d05713..01fcdf7 100644
--- a/TestON/tests/USECASE/SegmentRouting/dependencies/Testcaselib.py
+++ b/TestON/tests/USECASE/SegmentRouting/dependencies/Testcaselib.py
@@ -1071,3 +1071,21 @@
             Testcaselib.saveOnosDiagnostics( main )
             Testcaselib.cleanup( main, copyKarafLog=False, removeHostComponent=True )
             main.skipCase()
+
+    @staticmethod
+    def verifyHostLocation( main, hostName, locations, ipv6=False ):
+        """
+        Verify if the specified host is discovered by ONOS on the given locations
+        Required:
+            hostName: name of the host. ex. "h1"
+            locations: expected locations of the host. ex. "of:0000000000000005/8"
+                       Could be a string or list
+        Optional:
+            ipv6: Use True for IPv6 only hosts
+        Returns:
+            main.TRUE if host is discovered on all locations provided, otherwise main.FALSE
+        """
+        main.step( "Verify host {} is discovered at {}".format( hostName, locations ) )
+        hostIp = main.Network.getIPAddress( hostName, proto='IPV6' if ipv6 else 'IPV4' )
+        result = main.Cluster.active( 0 ).CLI.verifyHostLocation( hostIp, locations )
+        return result