[AutoTests] Update Chaos Tests to use APs & PCAPs from DataStore

As autotest is being depricated, this CL eleminates dependency on autotest.
Chaos tests now use autotest db to find and lock available APs (access points)
and PCAPs (Packet Capturers). The change proposed here will look for APs and
PCAPs in datastore that is available at appspot and is accessible via
https://chaos-188802.appspot.com/ . Database here is used to maintain status of
PCAPs & APs while test runs. During test run, test shall check if AP/PCAP is in
locked status and will lock it for test usage when found with unlocked status.
Once test completes using AP/PCAP, it shall then unlock the same.

BUG=chromium:1013809
TEST=Tested on Chaos-lab using test_that locally.

Change-Id: Ifcc99f675e26a7525ff933a8fcc755f966f3491f
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/1881841
Reviewed-by: Harpreet Grewal <harpreet@chromium.org>
Tested-by: Dinesh Kumar Sunkara <dsunkara@google.com>
Auto-Submit: Dinesh Kumar Sunkara <dsunkara@google.com>
Commit-Queue: Dinesh Kumar Sunkara <dsunkara@google.com>
diff --git a/server/cros/ap_configurators/ap_batch_locker.py b/server/cros/ap_configurators/ap_batch_locker.py
index 815989d..65d9a44 100644
--- a/server/cros/ap_configurators/ap_batch_locker.py
+++ b/server/cros/ap_configurators/ap_batch_locker.py
@@ -4,21 +4,20 @@
 
 import logging
 import random
-import requests
 
 from time import sleep
 
 import common
+from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib import utils
 from autotest_lib.server.cros.ap_configurators import \
     ap_configurator_factory
 from autotest_lib.client.common_lib.cros.network import ap_constants
 from autotest_lib.server.cros.ap_configurators import ap_cartridge
-
+from autotest_lib.server.cros.chaos_lib import chaos_datastore_utils as dutils
 
 # Max number of retry attempts to lock an ap.
 MAX_RETRIES = 3
-CHAOS_URL = 'https://chaos-188802.appspot.com'
 
 
 class ApLocker(object):
@@ -151,34 +150,53 @@
 
         return False
 
+
     def lock_ap_in_datastore(self, ap_locker):
-        """Locks an AP host in datastore.
+        """
+        Lock an AP host in datastore.
+
+        Test iterates through list of APs available in chaos_ap_list.conf file.
+        If AP with host_name is not found in datastore, this method adds the
+        same and shall use it for locking and testing.
 
         @param ap_locker: an ApLocker object, AP to be locked.
         @return a boolean, True iff ap_locker is locked.
+
         """
+        #ToDo dsunkara@: Find if below check is needed when not using Autotest.
         if not utils.host_is_in_lab_zone(ap_locker.configurator.host_name):
             ap_locker.to_be_locked = False
             return True
 
-        # Begin locking device in datastore.
-        locked_device = requests.put(CHAOS_URL + '/devices/lock', \
-                        json={"hostname":[ap_locker.configurator.host_name], \
-                        "locked_by":"TestRun"})
-        if locked_device.json()['result']:
-            self._locked_aps.append(ap_locker)
-            logging.info('locked %s', ap_locker.configurator.host_name)
+        # Get status of AP in datastore.
+        ap_device = dutils.show_device(ap_locker.configurator.host_name)
+        # Check if both Find / Add operations failed.
+        if not ap_device:
+            logging.error("Unable to find: %s in Datastore.",
+                          ap_locker.configurator.host_name)
             ap_locker.to_be_locked = False
-            return True
+        # Check lock status of AP before trying to lock it.
+        # Lock AP if its not locked, else retry
+        elif ap_device['lock_status']:
+            logging.error("AP is already locked by %s at %s",
+                          ap_device['locked_by'],
+                          ap_device['lock_status_updated'])
+            ap_locker.to_be_locked = False
         else:
-            ap_locker.retries -= 1
-            logging.info('%d retries left for %s',
+            # Lock device in datastore.
+            if dutils.lock_device(ap_locker.configurator.host_name):
+                self._locked_aps.append(ap_locker)
+                ap_locker.to_be_locked = False
+                return True
+            else:
+                ap_locker.retries -= 1
+                logging.info('%d retries left for %s',
                          ap_locker.retries,
                          ap_locker.configurator.host_name)
-            if ap_locker.retries == 0:
-                logging.info('No more retries left. Remove %s from list',
-                             ap_locker.configurator.host_name)
-                ap_locker.to_be_locked = False
+                if ap_locker.retries == 0:
+                    logging.info('No more retries left. Remove %s from list',
+                                 ap_locker.configurator.host_name)
+                    ap_locker.to_be_locked = False
 
         return False
 
@@ -199,9 +217,8 @@
 
             for ap_locker in self.aps_to_lock:
                 logging.info('checking %s', ap_locker.configurator.host_name)
-                # TODO(@rjahagir): Change method to datastore.
-                # if self.lock_ap_in_datastore(ap_locker):
-                if self.lock_ap_in_afe(ap_locker):
+                # Lock AP in DataStore
+                if self.lock_ap_in_datastore(ap_locker):
                     ap_batch.append(ap_locker.configurator)
                     if len(ap_batch) == batch_size:
                         break
@@ -243,38 +260,39 @@
         logging.error('Tried to unlock a host we have not locked (%s)?',
                       host_name)
 
-    def unlock_one_ap_in_datastore(self, host_name):
-        """Unlock one AP from datastore after we're done.
 
-        @param host_name: a string, host name.
+    def unlock_one_ap_in_datastore(self, host_name):
+        """
+        Unlock one AP from datastore after we're done.
+
+        @param host_name: a string, AP host name.
+
+        @raise TestError: when unable to unlock AP in datastore.
+
         """
         for ap_locker in self._locked_aps:
             if host_name == ap_locker.configurator.host_name:
                 # Unlock in datastore
-                unlocked_device = requests.put(CHAOS_URL + '/devices/unlock', \
-                                  json={"hostname":host_name})
-                # TODO: Raise error if unable to unlock.
-                if not unlocked_device.json()['result']:
-                    raise error
-                    logging.debug(unlocked_device.content())
+                # ToDo @dsunkara: change method name to unlock_one_ap
+                # once all dependencies are cleared on AFE.
+                unlocked_device = dutils.unlock_device(host_name)
+                if not unlocked_device:
+                    raise error.TestError('Failed to unlock AP: %s',
+                                          host_name)
                 else:
                     self._locked_aps.remove(ap_locker)
                 return
 
-        logging.error('Tried to unlock a host we have not locked (%s)?',
-                      host_name)
-
 
     def unlock_aps(self):
         """Unlock APs after we're done."""
         # Make a copy of all of the hostnames to process
+
         host_names = list()
         for ap_locker in self._locked_aps:
             host_names.append(ap_locker.configurator.host_name)
         for host_name in host_names:
-            # TODO(@rjahagir): Change method to datastore.
-            # self.unlock_one_ap_in_datastore(host_name)
-            self.unlock_one_ap(host_name)
+            self.unlock_one_ap_in_datastore(host_name)
 
 
     def unlock_and_reclaim_ap(self, host_name):
@@ -285,9 +303,7 @@
         for ap_locker in self._locked_aps:
             if host_name == ap_locker.configurator.host_name:
                 self.aps_to_lock.append(ap_locker)
-                # TODO(@rjahagir): Change method to datastore.
-                # self.unlock_one_ap_in_datastore(host_name)
-                self.unlock_one_ap(host_name)
+                self.unlock_one_ap_in_datastore(host_name)
                 return
 
 
diff --git a/server/cros/chaos_lib/chaos_datastore_utils.py b/server/cros/chaos_lib/chaos_datastore_utils.py
new file mode 100644
index 0000000..bbdd865
--- /dev/null
+++ b/server/cros/chaos_lib/chaos_datastore_utils.py
@@ -0,0 +1,224 @@
+# -*- coding: utf-8 -*-
+#!/usr/bin/env python2.7
+# Copyright (c) 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import json
+import logging
+import requests
+
+
+"""This file consists of all the helper methods needed to interact with the
+   Datastore @ https://chaos-188802.appspot.com/ used for ChromeOS Interop
+   testing.
+"""
+
+CHAOS_DATASTORE_URL = "https://chaos-188802.appspot.com"
+
+# The Datastore defines the following paths for operating methods.
+ADD_DEVICE = "devices/new"
+REMOVE_DEVICE = "devices/delete"
+LOCK_DEVICE = "devices/lock"
+UNLOCK_DEVICE = "devices/unlock"
+SHOW_DEVICE = "devices/"
+GET_DEVICES = "devices/"
+GET_UNLOCKED_DEVICES = "unlocked_devices/"
+GET_DEVICES_BY_AP_LABEL = "devices/location"
+
+# HTTP content type. JSON encoded with UTF-8 character encoding.
+HTTP_HEADER = {'content-type': 'application/json'}
+
+
+def add_device(host_name, ap_label):
+    """
+    Add a device(AP or Packet Capturer) in datastore.
+
+    @param host_name: string, hostname of the device.
+    @param ap_label: string, CrOS_AP (for AP), CrOS_PCAP (for PCAP)
+    @param lab_label: string, CrOS_Chaos (lab name), used for all APs & PCAPs
+
+    @return: True if device was added successfully; False otherwise.
+    @rtype: bool
+
+    """
+    request = CHAOS_DATASTORE_URL + '/' + ADD_DEVICE
+    logging.debug("Request = %s", request)
+    response = requests.post(request,
+                             headers=HTTP_HEADER,
+                             data=json.dumps({"hostname":host_name,
+                                              "ap_label":ap_label,
+                                              "lab_label":"CrOS_Chaos",
+                                              "router_name":host_name}))
+    if response.json()['result']:
+        logging.info("Added device %s to datastore", host_name)
+        return True
+
+    return False
+
+
+def remove_device(host_name):
+    """
+    Delete a device(AP or Packet Capturer) in datastore.
+
+    @param host_name: string, hostname of the device to delete.
+
+    @return: True if device was deleted successfully; False otherwise.
+    @rtype: bool
+
+    """
+    request = CHAOS_DATASTORE_URL + '/' + REMOVE_DEVICE
+    logging.debug("Request = %s", request)
+    response = requests.put(request,
+                            headers=HTTP_HEADER,
+                            data=json.dumps({"hostname":host_name}))
+    result_str = "%s deleted." % host_name
+    if result_str in response.text:
+        logging.info("Removed device %s from datastore", host_name)
+        return True
+
+    return False
+
+
+def lock_device(host_name):
+    """
+    Lock a device(AP or Packet Capturer) in datastore.
+
+    @param host_name: string, hostname of the device in datastore.
+
+    @return: True if operation was successful; False otherwise.
+    @rtype: bool
+
+    """
+    request = CHAOS_DATASTORE_URL + '/' + LOCK_DEVICE
+    logging.debug("Request = %s", request)
+    response = requests.put(request,
+                            headers=HTTP_HEADER,
+                            data=json.dumps({"hostname":host_name,
+                                             "locked_by":"TestRun"}))
+    if response.json()['result']:
+        logging.info("Locked device %s in datastore", host_name)
+        return True
+
+    return False
+
+
+def unlock_device(host_name):
+    """
+    Un-lock a device(AP or Packet Capturer) in datastore.
+
+    @param host_name: string, hostname of the device in datastore.
+
+    @return: True if operation was successful; False otherwise.
+    @rtype: bool
+
+    """
+    request = CHAOS_DATASTORE_URL + '/' + UNLOCK_DEVICE
+    logging.debug("Request = %s", request)
+    response = requests.put(request,
+                            headers=HTTP_HEADER,
+                            data=json.dumps({"hostname":host_name}))
+    if response.json()['result']:
+        logging.info("Finished un-locking AP %s in datastore", host_name)
+        return True
+
+    logging.error("Unable to unlock AP %s", host_name)
+    return False
+
+
+def show_device(host_name):
+    """
+    Show device properties for a given device(AP or Packet Capturer).
+
+    @param host_name: string, hostname of the device in datastore to fetch info.
+
+    @return: dict of device name:value properties if successful;
+             False otherwise.
+    @rtype: dict when True, else bool:False
+
+    """
+    request = CHAOS_DATASTORE_URL + '/' + SHOW_DEVICE + host_name
+    logging.debug("Request = %s", request)
+    response = requests.get(request)
+    if 'error' in response.text:
+        return False
+
+    return response.json()
+
+
+def get_unlocked_devices():
+    """
+    Get a list of all un-locked devices in the datastore.
+
+    @return: dict of all un-locked devices' name:value properties if successful;
+             False otherwise.
+    @rtype: dict when True, else bool:False
+
+    """
+    request = CHAOS_DATASTORE_URL + '/' + GET_UNLOCKED_DEVICES
+    logging.debug("Request = %s", request)
+    response = requests.get(request)
+    if 'error' in response.text:
+        return False
+
+    return response.json()
+
+
+def get_devices():
+    """
+    Get a list of all devices in the datastore.
+
+    @return: dict of all devices' name:value properties if successful;
+             False otherwise.
+    @rtype: dict when True, else bool:False
+
+    """
+    request = CHAOS_DATASTORE_URL + '/' + GET_DEVICES
+    logging.debug("Request = %s", request)
+    response = requests.get(request)
+    if 'error' in response.text:
+        return False
+
+    return response.json()
+
+
+def get_devices_by_type(ap_label, lab_label):
+    """
+    Get list of all un-locked devices by ap_label & lab_label
+
+    @param ap_label: string, CrOS_AP/CrOS_PCAP, to filter device types.
+    @param lab_label: string, "CrOS_Chaos", All devices in ChromeOS Chaos lab
+
+    @return: dict of all devices' name:value properties if successful;
+             False otherwise.
+    @rtype: dict when True, else bool:False
+
+    """
+    request = CHAOS_DATASTORE_URL + '/' +  GET_DEVICES_BY_AP_LABEL
+    logging.debug("Request = %s", request)
+    response = requests.put(request,
+                            headers=HTTP_HEADER,
+                            data=json.dumps({"ap_label":ap_label,
+                                             "lab_label":lab_label}))
+    if 'error' in response.text:
+        return False
+
+    return response.json()
+
+
+def find_device(host_name):
+    """
+    Find if given device(AP or Packet Capturer) in DataStore.
+
+    @param host_name: string, hostname of the device in datastore to fetch info.
+    @return: True if found; False otherwise.
+    @rtype: bool
+
+    """
+    request = CHAOS_DATASTORE_URL + '/' + SHOW_DEVICE + host_name
+    logging.debug("Request = %s", request)
+    response = requests.get(request)
+    if 'null' in response.text:
+        return False
+
+    return True
diff --git a/server/cros/chaos_lib/static_runner.py b/server/cros/chaos_lib/static_runner.py
index 5b26218..2f3550c 100644
--- a/server/cros/chaos_lib/static_runner.py
+++ b/server/cros/chaos_lib/static_runner.py
@@ -63,12 +63,11 @@
         """
 
         lock_manager = host_lock_manager.HostLockManager()
-        host_prefix = self._host.hostname.split('-')[0]
 
         with host_lock_manager.HostsLockedBy(lock_manager):
-            capture_host = utils.allocate_packet_capturer(
-                    lock_manager, hostname=capturer_hostname,
-                    prefix=host_prefix)
+            capture_host = utils.allocate_packet_capturer_in_datastore(
+                    lock_manager)
+
             # Cleanup and reboot packet capturer before the test.
             utils.sanitize_client(capture_host)
             capturer = site_linux_system.LinuxSystem(capture_host, {},
@@ -101,7 +100,7 @@
                         'chaos test?!')
 
             if conn_worker is not None:
-                work_client_machine = utils.allocate_packet_capturer(
+                work_client_machine = utils.allocate_packet_capturer_in_datastore(
                         lock_manager, hostname=work_client_hostname)
                 conn_worker.prepare_work_client(work_client_machine)
 
diff --git a/server/cros/host_lock_manager.py b/server/cros/host_lock_manager.py
index 2c88c20..39001f5 100644
--- a/server/cros/host_lock_manager.py
+++ b/server/cros/host_lock_manager.py
@@ -7,6 +7,7 @@
 import common
 
 from autotest_lib.server import site_utils
+from autotest_lib.server.cros.chaos_lib import chaos_datastore_utils as dutils
 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
 
 """HostLockManager class, for the dynamic_suite module.
@@ -99,6 +100,40 @@
         return mod_host
 
 
+    def lock_pcap_in_datastore(self, pcap_name):
+        """Lock Packet Capturere to use for the test in datastore."""
+
+        updated_hosts = set()
+        # Lock PCAP for capture
+        if dutils.lock_device(pcap_name):
+            logging.info("Locked Packet Capture device: %s", pcap_name)
+            updated_hosts.add(pcap_name)
+            self._locked_hosts = self._locked_hosts.union(updated_hosts)
+            return True
+
+        logging.error("Failed to lock %s PCAP.", pcap_name)
+        return False
+
+
+    def unlock_pcap_in_datastore(self, pcap_name=None):
+        """Unlock Packet Capturere in datastore after use."""
+
+        updated_hosts = self._locked_hosts
+        if not updated_hosts:
+            return False
+        logging.info('Unlocking pcap_host: %s', updated_hosts)
+        # Un-Lock PCAP
+        for pcap_host in updated_hosts:
+            if dutils.unlock_device(pcap_host):
+                logging.info("Locked Packet Capture device: %s", pcap_host)
+                updated_hosts.add(pcap_host)
+                self._locked_hosts = self._locked_hosts.intersection(updated_hosts)
+                return True
+
+        logging.error("Failed to lock %s PCAP.", pcap_name)
+        return False
+
+
     def lock(self, hosts, lock_reason='Locked by HostLockManager'):
         """Attempt to lock hosts in AFE.
 
@@ -201,4 +236,6 @@
 
     def __exit__(self, exntype, exnvalue, backtrace):
         signal.signal(signal.SIGTERM, self._old_handler)
+        # ToDO: dsunkara@ Cleanup methods if not using AutoTest
         self._manager.unlock()
+        self._manager.unlock_pcap_in_datastore()
diff --git a/server/cros/network/chaos_clique_utils.py b/server/cros/network/chaos_clique_utils.py
index 5536b4c..6403eae 100644
--- a/server/cros/network/chaos_clique_utils.py
+++ b/server/cros/network/chaos_clique_utils.py
@@ -19,6 +19,7 @@
 from autotest_lib.server.cros.ap_configurators import ap_configurator
 from autotest_lib.server.cros.ap_configurators import ap_cartridge
 from autotest_lib.server.cros.ap_configurators import ap_spec as ap_spec_module
+from autotest_lib.server.cros.chaos_lib import chaos_datastore_utils as dutils
 
 
 def allocate_packet_capturer(lock_manager, hostname, prefix):
@@ -51,6 +52,31 @@
     raise error.TestError('Unable to lock any pcaps - check in cautotest if '
                           'pcaps in %s are locked.', prefix)
 
+
+def allocate_packet_capturer_in_datastore(lock_manager):
+    """Finds a packet capturer to capture packets.
+
+    Locks the allocated pcap if it is discovered in datastore
+
+    @param lock_manager HostLockManager object.
+
+    @return: An SSHHost object representing a locked packet_capture machine.
+    """
+    # Gets available PCAPs that are NOT locked
+    available_pcaps = dutils.get_devices_by_type(ap_label='CrOS_PCAP',
+                                                 lab_label='CrOS_Chaos')
+    for pcap in available_pcaps:
+        # Ensure the pcap and dut are in the same subnet
+        if lock_manager.lock_pcap_in_datastore(pcap['hostname']):
+            return hosts.SSHHost(pcap['hostname'] + '.cros')
+            break
+        else:
+            logging.info('Unable to lock %s', pcap['hostname'])
+            continue
+    raise error.TestError('Unable to lock any pcaps - check datastore for '
+                          'pcaps locked status')
+
+
 def allocate_webdriver_instance(lock_manager):
     """Allocates a machine to capture webdriver instance.
 
@@ -202,7 +228,7 @@
     except ap_configurator.PduNotResponding as e:
         if ap.pdu not in broken_pdus:
             broken_pdus.append(ap.pdu)
-    batch_locker.unlock_one_ap(ap.host_name)
+    batch_locker.unlock_one_ap_in_datastore(ap.host_name)
 
 
 def filter_quarantined_and_config_failed_aps(aps, batch_locker, job,
@@ -240,7 +266,7 @@
                 release_ap(ap, batch_locker, broken_pdus)
             else:
                 # Cannot use _release_ap, since power_down will fail
-                batch_locker.unlock_one_ap(ap.host_name)
+                batch_locker.unlock_one_ap_in_datastore(ap.host_name)
     return list(set(aps) - set(aps_to_remove))