Randomly select TEE instead of always picking first one

src/opengradient/client/tee_registry.py

@@ -1,6 +1,7 @@
 """TEE Registry client for fetching verified TEE endpoints and TLS certificates."""
 
 import logging
+import random
 import ssl
 from dataclasses import dataclass
 from typing import List, NamedTuple, Optional
@@ -109,7 +110,7 @@ def get_active_tees_by_type(self, tee_type: int) -> List[TEEEndpoint]:
 
     def get_llm_tee(self) -> Optional[TEEEndpoint]:
         """
-        Return the first active LLM proxy TEE from the registry.
+        Return a random active LLM proxy TEE from the registry.
 
         Returns:
             TEEEndpoint for an active LLM proxy TEE, or None if none are available.
@@ -119,7 +120,7 @@ def get_llm_tee(self) -> Optional[TEEEndpoint]:
             logger.warning("No active LLM proxy TEEs found in registry")
             return None
 
-        return tees[0]
+        return random.choice(tees)
 
 
 def build_ssl_context_from_der(der_cert: bytes) -> ssl.SSLContext:

tests/tee_registry_test.py

@@ -144,7 +144,7 @@ def test_validator_type_label(self, mock_contract):
 
 
 class TestGetLlmTee:
-    def test_returns_first_active_tee(self, mock_contract):
+    def test_returns_random_active_tee(self, mock_contract):
         registry, contract = mock_contract
 
         contract.functions.getActiveTEEs.return_value.call.return_value = [
@@ -155,7 +155,7 @@ def test_returns_first_active_tee(self, mock_contract):
         result = registry.get_llm_tee()
 
         assert result is not None
-        assert result.endpoint == "https://tee-1.example.com"
+        assert result.endpoint in ("https://tee-1.example.com", "https://tee-2.example.com")
 
     def test_returns_none_when_no_tees(self, mock_contract):
         registry, contract = mock_contract