in memory cert + registry test

Author: balogh.adam@icloud.com

Makefile

@@ -31,7 +31,7 @@ docs:
 # Testing
 # ============================================================================
 
-test: utils_test client_test langchain_adapter_test opg_token_test
+test: utils_test client_test langchain_adapter_test opg_token_test tee_registry_test
 
 utils_test:
 	pytest tests/utils_test.py -v
@@ -45,6 +45,9 @@ langchain_adapter_test:
 opg_token_test:
 	pytest tests/opg_token_test.py -v
 
+tee_registry_test:
+	pytest tests/tee_registry_test.py -v
+
 integrationtest:
 	python integrationtest/agent/test_agent.py
 	python integrationtest/workflow_models/test_workflow_models.py

src/opengradient/client/tee_registry.py

@@ -2,7 +2,6 @@
 
 import logging
 import ssl
-import tempfile
 from dataclasses import dataclass
 from typing import List, Optional
 
@@ -138,15 +137,8 @@ def build_ssl_context_from_der(der_cert: bytes) -> ssl.SSLContext:
     """
     pem = ssl.DER_cert_to_PEM_cert(der_cert)
 
-    cert_file = tempfile.NamedTemporaryFile(
-        prefix="og_tee_tls_", suffix=".pem", delete=False, mode="w"
-    )
-    cert_file.write(pem)
-    cert_file.flush()
-    cert_file.close()
-
     ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
-    ctx.load_verify_locations(cert_file.name)
+    ctx.load_verify_locations(cadata=pem)
     ctx.check_hostname = False  # TEE cert may be issued for a hostname; we connect via IP
     ctx.verify_mode = ssl.CERT_REQUIRED
     return ctx

tests/langchain_adapter_test.py

@@ -26,34 +26,34 @@ def mock_client():
 @pytest.fixture
 def model(mock_client):
     """Create an OpenGradientChatModel with a mocked client."""
-    return OpenGradientChatModel(private_key="0x" + "a" * 64, model_cid=TEE_LLM.GPT_4O)
+    return OpenGradientChatModel(private_key="0x" + "a" * 64, model_cid=TEE_LLM.GPT_5)
 
 
 class TestOpenGradientChatModel:
     def test_initialization(self, model):
         """Test model initializes with correct fields."""
-        assert model.model_cid == TEE_LLM.GPT_4O
+        assert model.model_cid == TEE_LLM.GPT_5
         assert model.max_tokens == 300
         assert model.x402_settlement_mode == x402SettlementMode.SETTLE_BATCH
         assert model._llm_type == "opengradient"
 
     def test_initialization_custom_max_tokens(self, mock_client):
         """Test model initializes with custom max_tokens."""
-        model = OpenGradientChatModel(private_key="0x" + "a" * 64, model_cid=TEE_LLM.CLAUDE_3_5_HAIKU, max_tokens=1000)
+        model = OpenGradientChatModel(private_key="0x" + "a" * 64, model_cid=TEE_LLM.CLAUDE_HAIKU_4_5, max_tokens=1000)
         assert model.max_tokens == 1000
 
     def test_initialization_custom_settlement_mode(self, mock_client):
         """Test model initializes with custom settlement mode."""
         model = OpenGradientChatModel(
             private_key="0x" + "a" * 64,
-            model_cid=TEE_LLM.GPT_4O,
+            model_cid=TEE_LLM.GPT_5,
             x402_settlement_mode=x402SettlementMode.SETTLE,
         )
         assert model.x402_settlement_mode == x402SettlementMode.SETTLE
 
     def test_identifying_params(self, model):
         """Test _identifying_params returns model name."""
-        assert model._identifying_params == {"model_name": TEE_LLM.GPT_4O}
+        assert model._identifying_params == {"model_name": TEE_LLM.GPT_5}
 
 
 class TestGenerate:
@@ -210,7 +210,7 @@ def test_passes_correct_params_to_client(self, model, mock_client):
         model._generate([HumanMessage(content="Hi")], stop=["END"])
 
         mock_client.llm.chat.assert_called_once_with(
-            model=TEE_LLM.GPT_4O,
+            model=TEE_LLM.GPT_5,
             messages=[{"role": "user", "content": "Hi"}],
             stop_sequence=["END"],
             max_tokens=300,

tests/tee_registry_test.py

@@ -0,0 +1,250 @@
+import os
+import ssl
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
+
+from src.opengradient.client.tee_registry import (
+    TEE_TYPE_LLM_PROXY,
+    TEE_TYPE_VALIDATOR,
+    TEEEndpoint,
+    TEERegistry,
+    build_ssl_context_from_der,
+)
+
+
+# --- Helpers ---
+
+
+def _make_tee_info(
+    endpoint="https://tee.example.com",
+    payment_address="0xPayment",
+    tls_cert_der=b"\x01\x02\x03",
+    active=True,
+):
+    """Build a tuple matching the TEEInfo struct order from the contract."""
+    return (
+        "0xOwner",          # owner
+        payment_address,    # paymentAddress
+        endpoint,           # endpoint
+        b"pubkey",          # publicKey
+        tls_cert_der,       # tlsCertificate
+        b"pcrhash",         # pcrHash
+        0,                  # teeType
+        active,             # active
+        1000,               # registeredAt
+        2000,               # lastUpdatedAt
+    )
+
+
+def _make_self_signed_der() -> bytes:
+    """Generate a minimal self-signed DER certificate for testing."""
+    from cryptography import x509
+    from cryptography.hazmat.primitives import hashes, serialization
+    from cryptography.hazmat.primitives.asymmetric import rsa
+    from cryptography.x509.oid import NameOID
+    import datetime
+
+    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    subject = issuer = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "test")])
+    cert = (
+        x509.CertificateBuilder()
+        .subject_name(subject)
+        .issuer_name(issuer)
+        .public_key(key.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(datetime.datetime.now(datetime.UTC))
+        .not_valid_after(datetime.datetime.now(datetime.UTC) + datetime.timedelta(days=1))
+        .sign(key, hashes.SHA256())
+    )
+    return cert.public_bytes(serialization.Encoding.DER)
+
+
+# --- Fixtures ---
+
+
+@pytest.fixture
+def mock_contract():
+    """Create a TEERegistry with a mocked Web3 contract."""
+    with (
+        patch("src.opengradient.client.tee_registry.Web3") as mock_web3_cls,
+        patch("src.opengradient.client.tee_registry.get_abi") as mock_get_abi,
+    ):
+        mock_get_abi.return_value = []
+        mock_web3 = MagicMock()
+        mock_web3_cls.return_value = mock_web3
+        mock_web3_cls.HTTPProvider.return_value = MagicMock()
+        mock_web3_cls.to_checksum_address.side_effect = lambda x: x
+
+        contract = MagicMock()
+        mock_web3.eth.contract.return_value = contract
+
+        registry = TEERegistry(rpc_url="http://localhost:8545", registry_address="0xRegistry")
+        yield registry, contract
+
+
+# --- TEERegistry Tests ---
+
+
+class TestGetActiveTeesByType:
+    def test_returns_active_tees(self, mock_contract):
+        registry, contract = mock_contract
+
+        tee_id = b"\xaa" * 32
+        contract.functions.getTEEsByType.return_value.call.return_value = [tee_id]
+        contract.functions.getTEE.return_value.call.return_value = _make_tee_info()
+
+        result = registry.get_active_tees_by_type(TEE_TYPE_LLM_PROXY)
+
+        assert len(result) == 1
+        assert result[0].tee_id == tee_id.hex()
+        assert result[0].endpoint == "https://tee.example.com"
+        assert result[0].payment_address == "0xPayment"
+        assert result[0].tls_cert_der == b"\x01\x02\x03"
+
+    def test_skips_inactive_tees(self, mock_contract):
+        registry, contract = mock_contract
+
+        tee_id = b"\xbb" * 32
+        contract.functions.getTEEsByType.return_value.call.return_value = [tee_id]
+        contract.functions.getTEE.return_value.call.return_value = _make_tee_info(active=False)
+
+        result = registry.get_active_tees_by_type(TEE_TYPE_LLM_PROXY)
+        assert len(result) == 0
+
+    def test_skips_tee_with_empty_endpoint(self, mock_contract):
+        registry, contract = mock_contract
+
+        tee_id = b"\xcc" * 32
+        contract.functions.getTEEsByType.return_value.call.return_value = [tee_id]
+        contract.functions.getTEE.return_value.call.return_value = _make_tee_info(endpoint="")
+
+        result = registry.get_active_tees_by_type(TEE_TYPE_LLM_PROXY)
+        assert len(result) == 0
+
+    def test_skips_tee_with_empty_cert(self, mock_contract):
+        registry, contract = mock_contract
+
+        tee_id = b"\xdd" * 32
+        contract.functions.getTEEsByType.return_value.call.return_value = [tee_id]
+        contract.functions.getTEE.return_value.call.return_value = _make_tee_info(tls_cert_der=b"")
+
+        result = registry.get_active_tees_by_type(TEE_TYPE_LLM_PROXY)
+        assert len(result) == 0
+
+    def test_returns_empty_on_rpc_failure(self, mock_contract):
+        registry, contract = mock_contract
+
+        contract.functions.getTEEsByType.return_value.call.side_effect = Exception("RPC error")
+
+        result = registry.get_active_tees_by_type(TEE_TYPE_LLM_PROXY)
+        assert result == []
+
+    def test_skips_individual_tee_on_lookup_failure(self, mock_contract):
+        registry, contract = mock_contract
+
+        good_id = b"\xaa" * 32
+        bad_id = b"\xbb" * 32
+        contract.functions.getTEEsByType.return_value.call.return_value = [bad_id, good_id]
+
+        def get_tee_side_effect(tee_id):
+            mock = MagicMock()
+            if tee_id == bad_id:
+                mock.call.side_effect = Exception("lookup failed")
+            else:
+                mock.call.return_value = _make_tee_info()
+            return mock
+
+        contract.functions.getTEE.side_effect = get_tee_side_effect
+
+        result = registry.get_active_tees_by_type(TEE_TYPE_LLM_PROXY)
+        assert len(result) == 1
+        assert result[0].tee_id == good_id.hex()
+
+    def test_multiple_active_tees(self, mock_contract):
+        registry, contract = mock_contract
+
+        ids = [b"\x01" * 32, b"\x02" * 32, b"\x03" * 32]
+        contract.functions.getTEEsByType.return_value.call.return_value = ids
+
+        def get_tee_side_effect(tee_id):
+            mock = MagicMock()
+            mock.call.return_value = _make_tee_info(
+                endpoint=f"https://tee-{tee_id.hex()[:4]}.example.com"
+            )
+            return mock
+
+        contract.functions.getTEE.side_effect = get_tee_side_effect
+
+        result = registry.get_active_tees_by_type(TEE_TYPE_LLM_PROXY)
+        assert len(result) == 3
+
+    def test_validator_type_label(self, mock_contract):
+        """Ensure validator type queries work the same way."""
+        registry, contract = mock_contract
+
+        contract.functions.getTEEsByType.return_value.call.return_value = []
+
+        result = registry.get_active_tees_by_type(TEE_TYPE_VALIDATOR)
+        assert result == []
+        contract.functions.getTEEsByType.assert_called_once_with(TEE_TYPE_VALIDATOR)
+
+
+class TestGetLlmTee:
+    def test_returns_first_active_tee(self, mock_contract):
+        registry, contract = mock_contract
+
+        ids = [b"\x01" * 32, b"\x02" * 32]
+        contract.functions.getTEEsByType.return_value.call.return_value = ids
+        contract.functions.getTEE.return_value.call.return_value = _make_tee_info()
+
+        result = registry.get_llm_tee()
+
+        assert result is not None
+        assert result.tee_id == ids[0].hex()
+
+    def test_returns_none_when_no_tees(self, mock_contract):
+        registry, contract = mock_contract
+
+        contract.functions.getTEEsByType.return_value.call.return_value = []
+
+        result = registry.get_llm_tee()
+        assert result is None
+
+    def test_queries_llm_proxy_type(self, mock_contract):
+        registry, contract = mock_contract
+
+        contract.functions.getTEEsByType.return_value.call.return_value = []
+        registry.get_llm_tee()
+
+        contract.functions.getTEEsByType.assert_called_once_with(TEE_TYPE_LLM_PROXY)
+
+
+# --- build_ssl_context_from_der Tests ---
+
+
+class TestBuildSslContextFromDer:
+    def test_returns_ssl_context(self):
+        der_cert = _make_self_signed_der()
+        ctx = build_ssl_context_from_der(der_cert)
+
+        assert isinstance(ctx, ssl.SSLContext)
+
+    def test_hostname_check_disabled(self):
+        der_cert = _make_self_signed_der()
+        ctx = build_ssl_context_from_der(der_cert)
+
+        assert ctx.check_hostname is False
+
+    def test_cert_required(self):
+        der_cert = _make_self_signed_der()
+        ctx = build_ssl_context_from_der(der_cert)
+
+        assert ctx.verify_mode == ssl.CERT_REQUIRED
+
+    def test_rejects_invalid_der(self):
+        with pytest.raises(Exception):
+            build_ssl_context_from_der(b"not-a-valid-cert")