From 419818e6c10558fdb5ac1d0b1571448d0bf63dad Mon Sep 17 00:00:00 2001
From: darey-io <dare.oluumil@digital.justice.gov.uk>
Date: Mon, 15 Jan 2024 22:38:08 +0000
Subject: [PATCH] make ap versions configurable

---
 docker/Dockerfile    |  3 +++
 src/juniper.py       | 27 +++++++++++++++++++++------
 test/test_juniper.py | 15 +++++++++++++++
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index c07a078..be5fc23 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,5 +1,8 @@
 FROM python:3.11.5
 
+# Set a default value for AP_VERSIONS
+ENV AP_VERSIONS='{"AP45": "0.12.27139", "AP32": "0.12.27139"}'
+
 # Set the working directory to /app
 WORKDIR /app
 
diff --git a/src/juniper.py b/src/juniper.py
index 997ebb4..6b988ec 100644
--- a/src/juniper.py
+++ b/src/juniper.py
@@ -1,3 +1,5 @@
+import os
+import json
 import requests
 import json
 import getpass
@@ -9,6 +11,15 @@
 
 class Admin:
 
+    @staticmethod
+    def get_ap_versions():
+        ap_versions_str = os.getenv('AP_VERSIONS', '{}')
+        try:
+            return json.loads(ap_versions_str)
+        except json.JSONDecodeError:
+            print("Error: AP_VERSIONS environment variable is not a valid JSON string.")
+            return {}
+
     def login_via_username_and_password(self, username):
         # If no username defined ask user
         if username is None:
@@ -125,6 +136,9 @@ def build_payload(
         network_template_id,
         site_group_ids
 ):
+    # Get AP versions from environment variable
+    ap_versions = Admin.get_ap_versions()
+
     site = {'name': d.get('Site Name', ''),
             'address': d.get('Site Address', ''),
             "latlng": {"lat": d.get('gps', '')[0], "lng": d.get('gps', '')[1]},
@@ -145,10 +159,7 @@ def build_payload(
             "enabled": True,
             "version": "custom",
             "time_of_day": "02:00",
-            "custom_versions": {
-                "AP45": "0.12.27139",
-                "AP32": "0.12.27139"
-            },
+            "custom_versions": ap_versions,
             "day_of_week": ""
         },
 
@@ -331,11 +342,12 @@ def juniper_script(
         mist_login_method=None,
         site_group_ids=None,
         rf_template_id=None,
-        network_template_id=None
+        network_template_id=None,
+        ap_versions=None
 ):
     # Configure True/False to enable/disable additional logging of the API response objects
     show_more_details = True
-
+    ap_versions = Admin.get_ap_versions()
     # Check for required variables
     if org_id is None or org_id == '':
         raise ValueError('Please provide Mist org_id')
@@ -348,6 +360,9 @@ def juniper_script(
     if mist_login_method is None:
         print("mist_login_method not defined. Defaulting to credentials")
         mist_login_method = 'credentials'
+    if ap_versions is None or not isinstance(ap_versions, dict):
+            raise ValueError('Must provide a valid dictionary for ap_versions')
+
 
     # Prompt user if we are using production org_id
     warn_if_using_org_id_production(org_id)
diff --git a/test/test_juniper.py b/test/test_juniper.py
index c8339c9..65b0481 100644
--- a/test/test_juniper.py
+++ b/test/test_juniper.py
@@ -229,6 +229,21 @@ def test_put(self, mock_put, mock_successful_login, input_api_token):
         # Assert that the method returns the expected result
         self.assertIsNotNone(result)
 
+    @patch('src.juniper.Admin.get_ap_versions')
+    def test_ap_versions_handling(self, mock_get_ap_versions):
+        # Set up a valid AP_VERSIONS environment variable
+        valid_ap_versions = {"AP45": "0.12.27139", "AP32": "0.12.27139"}
+        mock_get_ap_versions.return_value = valid_ap_versions
+
+        # Test if juniper_script or Admin class handles valid AP_VERSIONS correctly
+        admin = Admin()
+        self.assertEqual(admin.get_ap_versions(), valid_ap_versions)
+
+        # Test with invalid AP_VERSIONS
+        with patch.dict('os.environ', {'AP_VERSIONS': 'invalid_json'}, clear=True):
+            with self.assertRaises(ValueError):
+                # Attempt to parse invalid JSON, should raise ValueError
+                admin.get_ap_versions()
 
 class TestCheckIfNeedToAppend(unittest.TestCase):