fix: resolve ruff linting errors and improve code quality

- Add missing return type annotations to example functions
- Fix exception handling with proper chaining (raise ... from e)
- Move inline imports to top-level in test files
- Add missing __init__.py file for test package structure
- Fix long line formatting issues throughout codebase
- Improve try/except/else structure consistency
- Update quote style consistency in string literals

Author: big14way

examples/03_turnkey_wallet_integration.py

@@ -15,6 +15,7 @@
 import asyncio
 import os
 from decimal import Decimal
+from typing import Any
 
 from flare_ai_kit.agents.turnkey import (
     AgentTransaction,
@@ -29,7 +30,7 @@
 from flare_ai_kit.wallet.turnkey_wallet import TurnkeySettings
 
 
-async def setup_turnkey_wallet():
+async def setup_turnkey_wallet() -> TurnkeyWallet:
     """Set up Turnkey wallet with proper configuration."""
     print("🔧 Setting up Turnkey wallet...")
 
@@ -78,7 +79,7 @@ async def setup_turnkey_wallet():
     return wallet
 
 
-async def create_demo_wallet(wallet: TurnkeyWallet):
+async def create_demo_wallet(wallet: TurnkeyWallet) -> tuple[str, str]:
     """Create a demo wallet for the AI agent."""
     print("\n💼 Creating demo wallet...")
 
@@ -99,7 +100,9 @@ async def create_demo_wallet(wallet: TurnkeyWallet):
         return "demo_wallet_id", "0x742d35Cc6634C0532925a3b8D8C8EE7c9e92bb1b"
 
 
-async def register_ai_agent(agent_connector: TurnkeyAgentConnector, wallet_id: str):
+async def register_ai_agent(
+    agent_connector: TurnkeyAgentConnector, wallet_id: str
+) -> AgentWalletConfig:
     """Register an AI agent with specific permissions."""
     print("\n🤖 Registering AI agent...")
 
@@ -127,7 +130,7 @@ async def register_ai_agent(agent_connector: TurnkeyAgentConnector, wallet_id: s
     return agent_config if success else None
 
 
-async def simulate_tee_operation(tee_manager: TEESecurityManager):
+async def simulate_tee_operation(_tee_manager: TEESecurityManager) -> str:
     """Simulate a TEE secure operation."""
     print("\n🔐 Simulating TEE secure operation...")
 
@@ -142,18 +145,13 @@ async def simulate_tee_operation(tee_manager: TEESecurityManager):
         # This would normally validate against real TEE attestation
         print("⚠️  Note: Using mock TEE attestation for demo purposes")
 
-        # Create secure operation (this will fail in demo without real TEE)
-        # secure_op = await tee_manager.create_secure_operation(
-        #     "ai_transaction",
-        #     operation_data,
-        #     mock_attestation_token
-        # )
+        # In production, this would create a secure operation
 
         print("✅ TEE operation simulated (would validate in production)")
         return mock_attestation_token
 
     except Exception as e:
-        print(f"⚠️  TEE validation would fail in production: {e}")
+        print(f"❌ TEE operation failed: {e}")
         print("   Using mock attestation for demo purposes")
         return mock_attestation_token
 
@@ -162,7 +160,7 @@ async def execute_ai_transaction(
     agent_connector: TurnkeyAgentConnector,
     agent_config: AgentWalletConfig,
     attestation_token: str,
-):
+) -> dict[str, Any]:
     """Execute a transaction initiated by the AI agent."""
     print("\n💸 Executing AI agent transaction...")
 
@@ -179,7 +177,9 @@ async def execute_ai_transaction(
     agent_transaction = AgentTransaction(
         agent_id=agent_config.agent_id,
         transaction_request=transaction_request,
-        justification="Automated yield optimization - moving funds to higher yield protocol",
+        justification=(
+            "Automated yield optimization - moving funds to higher yield protocol"
+        ),
         confidence_score=0.85,
         risk_assessment="Low risk - established DeFi protocol with 99.9% uptime",
     )
@@ -283,7 +283,8 @@ async def demonstrate_policy_enforcement(
     )
     if not result["success"]:
         print(
-            f"✅ Policy correctly blocked transaction without justification: {result['error']}"
+            "✅ Policy correctly blocked transaction without justification: "
+            f"{result['error']}"
         )
     else:
         print("❌ Policy should have blocked this transaction")
@@ -302,12 +303,10 @@ async def show_agent_status_and_history(
         print(f"   💼 Wallet ID: {status['wallet_id']}")
         print(f"   📊 Total transactions: {status['statistics']['total_transactions']}")
         print(f"   💰 Total value: {status['statistics']['total_value_eth']:.6f} ETH")
-        print(
-            f"   📈 Average confidence: {status['statistics']['avg_confidence_score']:.2f}"
-        )
-        print(
-            f"   🕒 Recent transactions (24h): {status['statistics']['recent_transactions_24h']}"
-        )
+        avg_confidence = status["statistics"]["avg_confidence_score"]
+        print(f"   📈 Average confidence: {avg_confidence:.2f}")
+        recent_24h = status["statistics"]["recent_transactions_24h"]
+        print(f"   🕒 Recent transactions (24h): {recent_24h}")
 
     # Get transaction history
     history = await agent_connector.get_transaction_history(agent_id=agent_id, limit=10)

src/flare_ai_kit/agents/turnkey.py

@@ -221,12 +221,6 @@ async def execute_agent_transaction(
                 tx_hash=signed_tx.transaction_hash,
             )
 
-            return {
-                "success": True,
-                "transaction_hash": signed_tx.transaction_hash,
-                "signed_transaction": signed_tx.signed_transaction,
-            }
-
         except PermissionError as e:
             logger.exception(
                 "agent_transaction_denied", agent_id=agent_id, error=str(e)
@@ -238,6 +232,12 @@ async def execute_agent_transaction(
                 "agent_transaction_failed", agent_id=agent_id, error=str(e)
             )
             return {"success": False, "error": f"Transaction failed: {e}"}
+        else:
+            return {
+                "success": True,
+                "transaction_hash": signed_tx.transaction_hash,
+                "signed_transaction": signed_tx.signed_transaction,
+            }
 
     async def _validate_agent_transaction(
         self, agent_transaction: AgentTransaction, config: AgentWalletConfig
@@ -258,7 +258,10 @@ async def _validate_agent_transaction(
         if agent_transaction.confidence_score < min_confidence:
             return {
                 "valid": False,
-                "reason": f"Agent confidence {agent_transaction.confidence_score} below threshold {min_confidence}",
+                "reason": (
+                    f"Agent confidence {agent_transaction.confidence_score} "
+                    f"below threshold {min_confidence}"
+                ),
             }
 
         # Validate justification is provided
@@ -272,7 +275,10 @@ async def _validate_agent_transaction(
         if tx_value > config.max_transaction_value:
             return {
                 "valid": False,
-                "reason": f"Transaction value {tx_value} ETH exceeds agent limit {config.max_transaction_value} ETH",
+                "reason": (
+                    f"Transaction value {tx_value} ETH exceeds agent limit "
+                    f"{config.max_transaction_value} ETH"
+                ),
             }
 
         # Additional risk-based validations could be added here

src/flare_ai_kit/wallet/permissions.py

@@ -293,7 +293,10 @@ def _check_rate_limit(
                 PolicyViolation(
                     policy_name=policy.name,
                     violation_type="rate_limit_transactions",
-                    description=f"Too many transactions in {window.duration_minutes} minutes. Limit: {window.max_transactions}",
+                    description=(
+                        f"Too many transactions in {window.duration_minutes} "
+                        f"minutes. Limit: {window.max_transactions}"
+                    ),
                     suggested_action=PolicyAction.DENY,
                 )
             )
@@ -307,7 +310,10 @@ def _check_rate_limit(
                     PolicyViolation(
                         policy_name=policy.name,
                         violation_type="rate_limit_value",
-                        description=f"Total value in {window.duration_minutes} minutes would exceed limit {window.max_value} ETH",
+                        description=(
+                            f"Total value in {window.duration_minutes} minutes "
+                            f"would exceed limit {window.max_value} ETH"
+                        ),
                         suggested_action=PolicyAction.DENY,
                     )
                 )

src/flare_ai_kit/wallet/tee_security.py

@@ -47,7 +47,8 @@ async def create_secure_operation(
         Create a secure operation with TEE attestation.
 
         Args:
-            operation_type: Type of operation (e.g., "wallet_creation", "transaction_signing")
+            operation_type: Type of operation (e.g., "wallet_creation",
+                "transaction_signing")
             operation_data: Operation-specific data
             attestation_token: TEE attestation token
 
@@ -126,14 +127,15 @@ async def verify_operation_integrity(self, operation: SecureOperation) -> bool:
         # Re-validate TEE attestation
         try:
             self.vtpm_validator.validate_token(operation.attestation_token)
-            return True
         except Exception as e:
             logger.exception(
                 "TEE attestation re-validation failed",
                 operation_id=operation.operation_id,
                 error=str(e),
             )
             return False
+        else:
+            return True
 
     async def encrypt_sensitive_data(
         self,
@@ -158,7 +160,7 @@ async def encrypt_sensitive_data(
             claims = self.vtpm_validator.validate_token(attestation_token)
         except Exception as e:
             msg = f"Invalid TEE context: {e}"
-            raise ValueError(msg)
+            raise ValueError(msg) from e
 
         # Derive encryption key from TEE-specific data
         key_material = self._derive_tee_key(claims, additional_data)

tests/unit/agents/__init__.py

@@ -9,7 +9,7 @@
     AgentWalletConfig,
     TurnkeyAgentConnector,
 )
-from flare_ai_kit.wallet.base import TransactionRequest
+from flare_ai_kit.wallet.base import SignedTransaction, TransactionRequest
 from flare_ai_kit.wallet.turnkey_wallet import TurnkeyWallet
 
 
@@ -124,7 +124,7 @@ def test_confidence_score_validation(self):
         assert transaction.confidence_score == 0.75
 
         # Invalid confidence scores should raise validation error
-        with pytest.raises(ValueError):
+        with pytest.raises(ValueError, match="Confidence score must be"):
             AgentTransaction(
                 agent_id="test_agent",
                 transaction_request=TransactionRequest(
@@ -183,7 +183,6 @@ async def test_execute_agent_transaction_success(
         await agent_connector.register_agent(agent_config)
 
         # Mock successful transaction signing
-        from flare_ai_kit.wallet.base import SignedTransaction
 
         signed_tx = SignedTransaction(
             transaction_hash="0x123abc...",
@@ -382,7 +381,6 @@ async def test_get_transaction_history(
         await agent_connector.register_agent(agent_config)
 
         # Mock successful transaction
-        from flare_ai_kit.wallet.base import SignedTransaction
 
         signed_tx = SignedTransaction(
             transaction_hash="0x123abc...",

tests/unit/agents/test_turnkey.py

@@ -3,6 +3,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from cryptography.hazmat.primitives.asymmetric import ec
 
 from flare_ai_kit.wallet.base import TransactionRequest, WalletAddress
 from flare_ai_kit.wallet.permissions import PermissionEngine, PolicyViolation
@@ -269,7 +270,6 @@ def test_sign_request(self, turnkey_wallet):
         test_body = '{"test": "data"}'
 
         # Mock private key operations
-        from cryptography.hazmat.primitives.asymmetric import ec
 
         with patch(
             "cryptography.hazmat.primitives.serialization.load_pem_private_key"