From a7ac00db3f5636137745e68854850a22b519ada1 Mon Sep 17 00:00:00 2001
From: Mark Phelps <209477+markphelps@users.noreply.github.com>
Date: Wed, 10 Jan 2024 15:12:29 -0500
Subject: [PATCH] feat: use auth strategy for python

---
 flipt-python/flipt/__init__.py                |  6 +++---
 flipt-python/flipt/authentication/__init__.py | 19 +++++++++++++++++++
 flipt-python/flipt/evaluation/__init__.py     | 10 ++++------
 flipt-python/flipt/evaluation/models.py       |  1 -
 flipt-python/tests/__init__.py                |  5 ++++-
 5 files changed, 30 insertions(+), 11 deletions(-)
 create mode 100644 flipt-python/flipt/authentication/__init__.py

diff --git a/flipt-python/flipt/__init__.py b/flipt-python/flipt/__init__.py
index c5473fc..556f93e 100644
--- a/flipt-python/flipt/__init__.py
+++ b/flipt-python/flipt/__init__.py
@@ -1,13 +1,13 @@
 import typing
 from .evaluation import Evaluation
+from .authentication import AuthenticationStrategy
 
 
 class FliptClient:
     def __init__(
         self,
         url: str = "http://localhost:8080",
-        client_token: typing.Optional[str] = None,
-        jwt_token: typing.Optional[str] = None,
         timeout: int = 60,
+        authentication: typing.Optional[AuthenticationStrategy] = None,
     ):
-        self.evaluation = Evaluation(url, client_token, jwt_token, timeout)
+        self.evaluation = Evaluation(url, timeout, authentication)
diff --git a/flipt-python/flipt/authentication/__init__.py b/flipt-python/flipt/authentication/__init__.py
new file mode 100644
index 0000000..9007c2f
--- /dev/null
+++ b/flipt-python/flipt/authentication/__init__.py
@@ -0,0 +1,19 @@
+class AuthenticationStrategy:
+    def authenticate(self, headers: dict):
+        raise NotImplementedError()
+
+
+class ClientTokenAuthentication(AuthenticationStrategy):
+    def __init__(self, token: str):
+        self.token = token
+
+    def authenticate(self, headers: dict):
+        headers["Authorization"] = f"Bearer {self.token}"
+
+
+class JWTAuthentication(AuthenticationStrategy):
+    def __init__(self, token: str):
+        self.token = token
+
+    def authenticate(self, headers: dict):
+        headers["Authorization"] = f"JWT {self.token}"
diff --git a/flipt-python/flipt/evaluation/__init__.py b/flipt-python/flipt/evaluation/__init__.py
index 02b263e..e3e2a24 100644
--- a/flipt-python/flipt/evaluation/__init__.py
+++ b/flipt-python/flipt/evaluation/__init__.py
@@ -8,23 +8,21 @@
     EvaluationRequest,
     VariantEvaluationResponse,
 )
+from ..authentication import AuthenticationStrategy
 
 
 class Evaluation:
     def __init__(
         self,
         url: str,
-        client_token: typing.Optional[str],
-        jwt_token: typing.Optional[str],
         timeout: int,
+        authentication: typing.Optional[AuthenticationStrategy] = None,
     ):
         self.url = url
         self.headers = {}
-        if client_token != None:
-            self.headers["Authorization"] = f"Bearer {client_token}"
-        if jwt_token != None:
-            self.headers["Authorization"] = f"JWT {jwt_token}"
         self.timeout = timeout
+        if authentication:
+            authentication.authenticate(self.headers)
 
     def variant(self, request: EvaluationRequest) -> VariantEvaluationResponse:
         response = httpx.post(
diff --git a/flipt-python/flipt/evaluation/models.py b/flipt-python/flipt/evaluation/models.py
index ef6517d..9674ed9 100644
--- a/flipt-python/flipt/evaluation/models.py
+++ b/flipt-python/flipt/evaluation/models.py
@@ -1,4 +1,3 @@
-import enum
 from pydantic import BaseModel, Field
 from typing import List, Optional
 
diff --git a/flipt-python/tests/__init__.py b/flipt-python/tests/__init__.py
index e16f403..79d32bb 100644
--- a/flipt-python/tests/__init__.py
+++ b/flipt-python/tests/__init__.py
@@ -2,6 +2,7 @@
 import unittest
 from flipt import FliptClient
 from flipt.evaluation import BatchEvaluationRequest, EvaluationRequest
+from flipt.authentication import ClientTokenAuthentication
 
 
 class TestFliptEvaluationClient(unittest.TestCase):
@@ -14,7 +15,9 @@ def setUp(self) -> None:
         if auth_token is None:
             raise Exception("FLIPT_AUTH_TOKEN not set")
 
-        self.flipt_client = FliptClient(url=flipt_url, client_token=auth_token)
+        self.flipt_client = FliptClient(
+            url=flipt_url, authentication=ClientTokenAuthentication(auth_token)
+        )
 
     def test_variant(self):
         variant = self.flipt_client.evaluation.variant(