From 2979b00cdb0e665f5969d3330bed9eceed5ef23e Mon Sep 17 00:00:00 2001
From: Matthew Stidham <stidmatt@gmail.com>
Date: Thu, 18 Nov 2021 15:26:14 -0800
Subject: [PATCH] lf_ap_auto_test: Fix lf_attenmod_test

Signed-off-by: Matthew Stidham <stidmatt@gmail.com>
---
 py-json/lf_attenmod.py | 43 ++++++++++++++++++++++++------------------
 1 file changed, 25 insertions(+), 18 deletions(-)

diff --git a/py-json/lf_attenmod.py b/py-json/lf_attenmod.py
index f9652453..b13eca5e 100644
--- a/py-json/lf_attenmod.py
+++ b/py-json/lf_attenmod.py
@@ -18,15 +18,15 @@ class ATTENUATORProfile(LFCliBase):
         super().__init__(lfclient_host, lfclient_port, debug_)
         self.lfclient_host = lfclient_host
         self.COMMANDS = ["show_attenuators", "set_attenuator"]
-        self.atten_serno = ""
-        self.atten_idx = ""
-        self.atten_val = ""
+        self.atten_serno = None
+        self.atten_idx = None
+        self.atten_val = None
         self.atten_data = {
             "shelf": 1,
             "resource": 1,
-            "serno": None,
-            "atten_idx": None,
-            "val": None,
+            "serno": self.atten_serno,
+            "atten_idx": self.atten_idx,
+            "val": self.atten_val,
             "mode": None,
             "pulse_width_us5": None,
             "pulse_interval_ms": None,
@@ -37,10 +37,10 @@ class ATTENUATORProfile(LFCliBase):
 
     def set_command_param(self, command_name, param_name, param_value):
         # we have to check what the param name is
-        if (command_name is None) or (command_name == ""):
-            return
-        if (param_name is None) or (param_name == ""):
-            return
+        if not command_name:
+            raise ValueError("Command Name is required")
+        if not param_name:
+            raise ValueError("Paramater is required")
         if command_name not in self.COMMANDS:
             raise ValueError("Command name name [%s] not defined in %s" % (command_name, self.COMMANDS))
         if command_name == "set_attenuator":
@@ -53,19 +53,27 @@ class ATTENUATORProfile(LFCliBase):
         if response is None:
             print(response)
             raise ValueError("Cannot find any endpoints")
-        else:
-            attenuator_resp = response["attenuators"]
+        elif 'attenuator' or 'attenuators' in response.keys():
+            try:
+                attenuator_resp = response["attenuators"]
+            except KeyError:
+                attenuator_resp = [response["attenuator"]]
             for attenuator in attenuator_resp:
                 for key, val in attenuator.items():
                     if key == "entity id":
                         serial_num = val.split(".")
                         print("Serial-num : %s" % serial_num[-1])
-                print("%s : %s" % (key, val))
-        print("\n")
+                    print("%s : %s" % (key, val))
+                    print('\n')
+
+        else:
+            raise ValueError('No attenuators in response')
 
     def create(self):
-        if type(self.atten_serno) == str or len(self.atten_idx) == 0 or type(self.atten_val) == str:
-            print("ERROR:  Must specify atten_serno, atten_idx, and atten_val when setting attenuator.\n")
+        if int(self.atten_val) > 955:
+            raise ValueError("Attenuation ddB value must be 955 or less")
+        if int(self.atten_idx) > 7:
+            raise ValueError("Attenuation ddB value must be 7 or less")
         print("Setting Attenuator...")
         self.set_command_param("set_attenuator", "serno", self.atten_serno)
         self.set_command_param("set_attenuator", "atten_idx", self.atten_idx)
@@ -73,7 +81,6 @@ class ATTENUATORProfile(LFCliBase):
         set_attenuators = LFRequest.LFRequest(self.lfclient_url + "/cli-json/set_attenuator", debug_=self.debug)
         set_attenuators.addPostData(self.atten_data)
         time.sleep(0.01)
-        json_response = set_attenuators.jsonPost(self.debug)
+        set_attenuators.jsonPost(self.debug)
         time.sleep(10)
         print("\n")
-