From deb19d78ae415ea32c941d2551c7c488fb9b941f Mon Sep 17 00:00:00 2001
From: simonwt <simon.little@wikitribune.com>
Date: Mon, 13 Apr 2026 22:26:42 +0100
Subject: [PATCH] Handle token method for saving token token to a file

---
 .gitignore                           |  2 ++
 src/trustcafeapiwrapper/apiclient.py | 25 +++++++++++++++++++++++--
 testing.py                           | 17 ++++-------------
 3 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/.gitignore b/.gitignore
index d5bbe91..b66b49d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -26,4 +26,6 @@ dist/*
 
 # Test Artifacts
 token_data.json
+token_data_alpha.json
+token_dataproduction.json
 test_output.json
\ No newline at end of file
diff --git a/src/trustcafeapiwrapper/apiclient.py b/src/trustcafeapiwrapper/apiclient.py
index ab25a19..658c862 100644
--- a/src/trustcafeapiwrapper/apiclient.py
+++ b/src/trustcafeapiwrapper/apiclient.py
@@ -1,7 +1,6 @@
-import requests
+import os,requests
 import simplejson as json
 from urllib.parse import urlencode
-
 from pydantic import (
     PrivateAttr, BaseModel, SecretStr, HttpUrl
 )
@@ -153,6 +152,28 @@ class APIClient(BaseModel):
         current_time = int(time.time())
         return self._access_token and current_time < self._access_token_timeout
     
+    def handle_token(self, token_data_path=None):
+        """
+        Handle access token retrieval and caching.
+        """
+        if token_data_path is None:
+            token_data_path = f"token_data_{self.environment}.json" 
+
+        if os.path.exists(token_data_path):
+            with open(token_data_path, "r") as f:
+                token_data = json.load(f)
+            self.set_token(token_data)
+        else:
+            if self.debug:
+                print(f"No token data file found at '{token_data_path}'. A new token will be obtained.")
+                raise FileNotFoundError(f"Token data file '{token_data_path}' not found. Please ensure the file exists or handle token retrieval appropriately.")
+            
+        # Get a new one if we don't have one or if the existing one is expired
+        if not self.is_token_valid():
+            tokendata = self.sign_in()
+            with open(token_data_path, "w") as f:
+                json.dump(tokendata, f, indent=2)
+
     def run_job(self, job_function, *args, **kwargs):
         """
         Utility method to run a job function with the API client as the first argument.
diff --git a/testing.py b/testing.py
index ef2be5f..72c8cc2 100644
--- a/testing.py
+++ b/testing.py
@@ -26,16 +26,7 @@ API = APIClient(
 # Keep a token cache to avoid unnecessary sign-ins during development. 
 # (In production, you'd handle this more robustly and securely)
 
-if os.path.exists("token_data.json"):
-    with open("token_data.json", "r") as f:
-        token_data = json.load(f)
-    API.set_token(token_data)
-
-# Get a new one if we don't have one or if the existing one is expired
-if not API.is_token_valid():
-    tokendata = API.sign_in()
-    with open("token_data.json", "w") as f:
-        json.dump(tokendata, f, indent=2)
+API.handle_token()  # This will load the token from file if it exists and is valid, or sign in to get a new one if not.
 
 '''
 END IMPORTANT BIT
@@ -50,8 +41,8 @@ def save_response(response):
 
 
 # print("-----------Get a user profile----------------")
-# profile = API.run_job('userprofile.get', "simon-little")
-# print(profile)
+profile = API.run_job('userprofile.get', "simon-little")
+print(profile)
 # print("-------------- Get a branch -----------------")
 # branch = API.run_job('branch.get', "music")
 # print(branch)
@@ -138,7 +129,7 @@ def save_response(response):
 #     x += 1
 
 # save_response(API.run_job('notification.listnotifications'))
-save_response(API.run_job('notification.markallasread'))
+# save_response(API.run_job('notification.markallasread'))
 
 # save_response(API.run_job('feed.cafefeed'))