import unittest
from unittest.mock import patch

class TestAgent(unittest.TestCase):
    def setUp(self):
        self.client = ArmorIQClient(
            api_key="ak_live_" + "a" * 64,
            user_id="test_user",
            agent_id="test_agent"
        )

    @patch('armoriq_sdk.client.ArmorIQClient.invoke')
    def test_successful_invocation(self, mock_invoke):
        # Mock successful response
        mock_invoke.return_value = {
            "success": True,
            "data": {"result": 42},
            "execution_time_ms": 100
        }

        result = self.client.invoke("math-mcp", "calculate", "token", {})

        self.assertTrue(result["success"])
        self.assertEqual(result["data"]["result"], 42)

    @patch('armoriq_sdk.client.ArmorIQClient.get_intent_token')
    def test_token_expiration_handling(self, mock_get_token):
        # First call returns expired token
        # Second call returns fresh token
        mock_get_token.side_effect = [
            {"token": "expired_token", "expires_at": 0},
            {"token": "fresh_token", "expires_at": 9999999999}
        ]

        # Test auto-refresh logic
        ...

import { ArmorIQClient, IntentToken, MCPInvocationResult } from '@armoriq/sdk';
import { jest, describe, it, expect, beforeEach } from '@jest/globals';

describe('ArmorIQ Agent Tests', () => {
  let client: ArmorIQClient;

  beforeEach(() => {
    // Mock environment variables
    process.env.ARMORIQ_API_KEY = 'ak_test_' + 'a'.repeat(64);
    process.env.USER_ID = 'test_user';
    process.env.AGENT_ID = 'test_agent';

    client = new ArmorIQClient({
      apiKey: process.env.ARMORIQ_API_KEY,
      userId: process.env.USER_ID,
      agentId: process.env.AGENT_ID,
      useProduction: false
    });
  });

  it('should successfully invoke an action', async () => {
    // Mock the invoke method
    const mockResult: MCPInvocationResult = {
      mcp: 'math-mcp',
      action: 'calculate',
      result: { value: 42 },
      status: 'success',
      executionTime: 0.1,
      verified: true,
      metadata: {}
    };

    jest.spyOn(client, 'invoke').mockResolvedValue(mockResult);

    const result = await client.invoke(
      'math-mcp',
      'calculate',
      {} as IntentToken,
      {}
    );

    expect(result.status).toBe('success');
    expect(result.result.value).toBe(42);
  });

  it('should handle token expiration', async () => {
    const mockToken: Partial<IntentToken> = {
      tokenId: 'test-token',
      expiresAt: Date.now() / 1000 + 3600  // 1 hour from now
    };

    jest.spyOn(client, 'getIntentToken').mockResolvedValue(mockToken as IntentToken);

    const token = await client.getIntentToken({} as any);
    
    expect(IntentToken.isExpired(token as IntentToken)).toBe(false);
    expect(IntentToken.timeUntilExpiry(token as IntentToken)).toBeGreaterThan(3500);
  });
});