From 1cdd6284efd065afcabd5529605c29b8d346dbef Mon Sep 17 00:00:00 2001
From: Daofeng Wu <dolphin.w.e@gmail.com>
Date: Wed, 25 Dec 2024 19:49:32 +0900
Subject: [PATCH] refactor(tool/email_organizer/test): cache auth records

---
 .../__test__/invoice_organizer.py             | 27 ++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/npiai/tools/email_organizer/__test__/invoice_organizer.py b/npiai/tools/email_organizer/__test__/invoice_organizer.py
index 5b914de..7bf0f39 100644
--- a/npiai/tools/email_organizer/__test__/invoice_organizer.py
+++ b/npiai/tools/email_organizer/__test__/invoice_organizer.py
@@ -4,21 +4,45 @@
 from npiai import Context
 from npiai.tools.outlook import Outlook
 from npiai.tools.email_organizer import EmailOrganizer
-from azure.identity import InteractiveBrowserCredential
+from azure.identity import (
+    InteractiveBrowserCredential,
+    TokenCachePersistenceOptions,
+    AuthenticationRecord,
+)
+
+token_cache = ".cache/outlook_token_cache.json"
 
 
 async def main():
+    # authenticate
+    authentication_record = None
+
+    if os.path.exists(token_cache):
+        with open(token_cache, "r") as f:
+            authentication_record = AuthenticationRecord.deserialize(f.read())
+
     creds = InteractiveBrowserCredential(
         client_id=os.environ.get("AZURE_CLIENT_ID", None),
         # tenant_id=os.environ.get("AZURE_TENANT_ID", None),
         tenant_id="common",
+        cache_persistence_options=TokenCachePersistenceOptions(),
+        authentication_record=authentication_record,
     )
 
+    record = creds.authenticate(scopes=["Mail.Read", "Mail.Send"])
+
+    os.makedirs(os.path.dirname(token_cache), exist_ok=True)
+
+    with open(token_cache, "w") as f:
+        f.write(record.serialize())
+
     async with EmailOrganizer(provider=Outlook(creds)) as tool:
+        # list emails
         email_list = [email async for email in tool.list_inbox_stream(limit=10)]
 
         print("Raw email list:", json.dumps(email_list, indent=4, ensure_ascii=False))
 
+        # filter invoice-like emails
         filtered_emails = []
 
         async for result in tool.filter_stream(
@@ -34,6 +58,7 @@ async def main():
                 f'Subject: {result["email"]["subject"]}, Matched: {result["matched"]}'
             )
 
+        # summarize invoice-like emails
         async for item in tool.summarize_stream(
             ctx=Context(),
             email_or_id_list=filtered_emails,