fix(tests): corrected extension import logic, added enhancementas as per suggestions

Author: maxmekiska

examples/memory/encrypted_session_example.py

@@ -7,9 +7,11 @@
 """
 
 import asyncio
+from typing import cast
 
 from agents import Agent, Runner, SQLiteSession
 from agents.extensions.memory import EncryptedSession
+from agents.extensions.memory.encrypt_session import EncryptedEnvelope
 
 
 async def main():
@@ -27,7 +29,7 @@ async def main():
     session = EncryptedSession(
         session_id=session_id,
         underlying_session=underlying_session,
-        encryption_key="my-secret-encryption-key",  # In production, use ENCRYPTION_KEY env var
+        encryption_key="my-secret-encryption-key",
         ttl=3600,  # 1 hour TTL for messages
     )
 
@@ -89,8 +91,10 @@ async def main():
     print("Raw encrypted items in underlying storage:")
     for i, item in enumerate(raw_items, 1):
         if isinstance(item, dict) and item.get("__enc__") == 1:
+            enc_item = cast(EncryptedEnvelope, item)
             print(
-                f"  {i}. Encrypted envelope: __enc__={item['__enc__']}, payload length={len(item.get('payload', ''))}"
+                f"  {i}. Encrypted envelope: __enc__={enc_item['__enc__']}, "
+                f"payload length={len(enc_item['payload'])}"
             )
         else:
             print(f"  {i}. Unencrypted item: {item}")

pyproject.toml

@@ -66,6 +66,8 @@ dev = [
     "eval-type-backport>=0.2.2",
     "fastapi >= 0.110.0, <1",
     "aiosqlite>=0.21.0",
+    "types-cryptography",
+    "cryptography>=45.0, <46",
 ]
 
 [tool.uv.workspace]

src/agents/extensions/memory/__init__.py

@@ -8,10 +8,38 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 from .encrypt_session import EncryptedSession  # noqa: F401
 from .sqlalchemy_session import SQLAlchemySession  # noqa: F401
 
 __all__: list[str] = [
     "EncryptedSession",
     "SQLAlchemySession",
 ]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "EncryptedSession":
+        try:
+            from .encrypt_session import EncryptedSession
+
+            return EncryptedSession
+        except ModuleNotFoundError as e:
+            raise ImportError(
+                "EncryptedSession requires the 'cryptography' extra. "
+                "Install it with: pip install openai-agents[encrypt]"
+            ) from e
+
+    if name == "SQLAlchemySession":
+        try:
+            from .sqlalchemy_session import SQLAlchemySession
+
+            return SQLAlchemySession
+        except ModuleNotFoundError as e:
+            raise ImportError(
+                "SQLAlchemySession requires the 'sqlalchemy' extra. "
+                "Install it with: pip install openai-agents[sqlalchemy]"
+            ) from e
+
+    raise AttributeError(f"module {__name__} has no attribute {name}")

src/agents/extensions/memory/encrypt_session.py

@@ -18,7 +18,7 @@
     session = EncryptedSession(
         session_id="user-123",
         underlying_session=underlying_session,
-        encryption_key="your-encryption-key",  # or use ENCRYPTION_KEY env var
+        encryption_key="your-encryption-key",
         ttl=600,  # 10 minutes
     )
 
@@ -29,7 +29,6 @@
 
 import base64
 import json
-import os
 from typing import Any, Literal, TypedDict, TypeGuard, cast
 
 from cryptography.fernet import Fernet, InvalidToken
@@ -55,7 +54,7 @@ def _ensure_fernet_key_bytes(master_key: str) -> bytes:
     Returns raw bytes suitable for HKDF input.
     """
     if not master_key:
-        raise ValueError("ENCRYPTION_KEY missing; required for EncryptedSession.")
+        raise ValueError("encryption_key not set; required for EncryptedSession.")
     try:
         key_bytes = base64.urlsafe_b64decode(master_key)
         if len(key_bytes) == 32:
@@ -113,7 +112,7 @@ def __init__(
         self,
         session_id: str,
         underlying_session: SessionABC,
-        encryption_key: str | None = os.getenv("ENCRYPTION_KEY", None),
+        encryption_key: str,
         ttl: int = 600,
     ):
         """
@@ -126,8 +125,7 @@ def __init__(
         self.session_id = session_id
         self.underlying_session = underlying_session
         self.ttl = ttl
-        if encryption_key is None:
-            raise ValueError("ENCRYPTION_KEY missing; required for EncryptedSession.")
+
         master = _ensure_fernet_key_bytes(encryption_key)
         self.cipher = _derive_session_fernet_key(master, session_id)
         self._kid = "hkdf-v1"

tests/extensions/memory/test_encrypt_session.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-import os
 import tempfile
 import time
 from pathlib import Path
-from unittest.mock import patch
 
 import pytest
+
+pytest.importorskip("cryptography")  # Skip tests if cryptography is not installed
+
 from cryptography.fernet import Fernet
 
 from agents import Agent, Runner, SQLiteSession, TResponseInputItem
@@ -27,7 +28,7 @@ def agent() -> Agent:
 @pytest.fixture
 def encryption_key() -> str:
     """Fixture for a valid Fernet encryption key."""
-    return Fernet.generate_key().decode("utf-8")
+    return str(Fernet.generate_key().decode("utf-8"))
 
 
 @pytest.fixture
@@ -229,36 +230,6 @@ async def test_encrypted_session_pop_mixed_expired_valid(
     underlying_session.close()
 
 
-async def test_encrypted_session_env_key(underlying_session: SQLiteSession):
-    """Test encryption key from environment variable."""
-    key = Fernet.generate_key().decode("utf-8")
-
-    with patch.dict(os.environ, {"ENCRYPTION_KEY": key}):
-        session = EncryptedSession(
-            session_id="test_session",
-            underlying_session=underlying_session,
-        )
-
-        await session.add_items([{"role": "user", "content": "Test"}])
-        items = await session.get_items()
-        assert len(items) == 1
-        assert items[0].get("content") == "Test"
-
-    underlying_session.close()
-
-
-async def test_encrypted_session_missing_key(underlying_session: SQLiteSession):
-    """Test error handling for missing encryption key."""
-    with pytest.raises(ValueError, match="ENCRYPTION_KEY missing"):
-        EncryptedSession(
-            session_id="test_session",
-            underlying_session=underlying_session,
-            encryption_key=None,
-        )
-
-    underlying_session.close()
-
-
 async def test_encrypted_session_raw_string_key(underlying_session: SQLiteSession):
     """Test using raw string as encryption key (not base64)."""
     session = EncryptedSession(