From 590fae14072aaabb1cb02a683f7927c124a02f20 Mon Sep 17 00:00:00 2001
From: simonwt <simon.little@wikitribune.com>
Date: Tue, 7 Apr 2026 22:32:13 +0100
Subject: [PATCH] changing environment variables

---
 src/trustcafeapiwrapper/apiclient.py | 32 ++++++++++++++++++++++++++--
 tests/apiclient.py                   | 19 +++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/src/trustcafeapiwrapper/apiclient.py b/src/trustcafeapiwrapper/apiclient.py
index 0fe78ef..6e31fcb 100644
--- a/src/trustcafeapiwrapper/apiclient.py
+++ b/src/trustcafeapiwrapper/apiclient.py
@@ -11,7 +11,17 @@ endpoints = {
         "auth": "https://oo0wks9pbi.execute-api.us-east-1.amazonaws.com/alpha/",
         "content": "https://w1yygdhayc.execute-api.us-east-1.amazonaws.com/alpha/",
         "megaphone": "https://opdhjaktnl.execute-api.us-east-1.amazonaws.com/alpha/",
-    }
+        "moderation": "https://sbb1xrqsf1.execute-api.us-east-1.amazonaws.com/alpha/",
+        "spider": "https://xuvz1oj1sk.execute-api.us-east-1.amazonaws.com/dev/"
+    },
+    "production": {
+        "audrey": "https://iiouau5d2k.execute-api.us-east-1.amazonaws.com/production/",
+        "auth": "https://e2we3nktl4.execute-api.us-east-1.amazonaws.com/production/",
+        "content": "https://32hho6rvg1.execute-api.us-east-1.amazonaws.com/production/",
+        "megaphone": "https://yfnamdpnc8.execute-api.us-east-1.amazonaws.com/production/",
+        "moderation": "https://egfpoo3mw3.execute-api.us-east-1.amazonaws.com/production/",
+        "spider": "https://coi5kvgypc.execute-api.us-east-1.amazonaws.com/production/"
+    },
 }
 class APIClient(BaseModel):
     # Internal State (Not passed in __init__)
@@ -26,7 +36,14 @@ class APIClient(BaseModel):
     _access_token: str = PrivateAttr(default="")
     _access_token_timeout: int = PrivateAttr(default=0)
     
-        
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.validate_environment()
+    
+    def validate_environment(self):
+        if self.environment not in endpoints:
+            raise ValueError(f"Environment '{self.environment}' is not valid. Must be one of: {list(endpoints.keys())}") 
+
     def make_request(self, method: str, endpoint: str, path: str, data: dict = None, authenticate: bool = True, query_params: dict = None) -> dict:
         
         # Make sure the endpoint is defined in the endpoints dictionary
@@ -112,6 +129,17 @@ class APIClient(BaseModel):
         self._access_token = token_data.get("access_token", "")
         self._access_token_timeout = token_data.get("access_token_timeout", 0)
  
+    def set_environment(self, environment: str) -> None:
+        """
+        Set the environment for the API client.
+
+        Args:
+            environment (str): The environment to set (e.g., "alpha", "production").
+        """
+        
+        self.validate_environment()
+        self.environment = environment
+
     def is_token_valid(self) -> bool:
         """
         Checks if the current access token is still valid based on the current time and the token's expiration time.
diff --git a/tests/apiclient.py b/tests/apiclient.py
index 81f24f0..e34b38a 100644
--- a/tests/apiclient.py
+++ b/tests/apiclient.py
@@ -46,6 +46,25 @@ class TestAPIClient(unittest.TestCase):
         
         result = self.api_client.run_job(mock_job, "value1")
         self.assertEqual(result, "Job executed with value1")
+
+
+    def test_set_production_env(self):
+        Prod_API = APIClient(
+            client_id="prod_client_id",
+            client_secret="prod_client_secret",
+            debug=False,
+            environment="production"
+        )
+        self.assertEqual(Prod_API.environment, "production")
+
+    def test_set_bad_env(self):
+        with self.assertRaises(ValueError):
+            APIClient(
+                client_id="test_client_id",
+                client_secret="test_client_secret",
+                debug=False,
+                environment="invalid_env"
+            )
         
     @patch('trustcafeapiwrapper.apiclient.requests.request')
     def test_make_request(self, mock_request):