paygent-sdk 1.0.0__py3-none-any.whl

tests/test_client.py

@@ -0,0 +1,277 @@
+"""
+Tests for the Paygent SDK client.
+"""
+
+import json
+import logging
+import unittest
+from unittest.mock import Mock, patch, MagicMock
+
+import requests
+
+from paygent_sdk.client import Client
+from paygent_sdk.models import UsageData, UsageDataWithStrings, APIRequest, ModelPricing, MODEL_PRICING
+
+
+class TestClient(unittest.TestCase):
+    """Test cases for the Client class."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.api_key = "test-api-key"
+        self.base_url = "https://api.paygent.com"
+        self.client = Client(self.api_key, self.base_url)
+
+    def test_new_client(self):
+        """Test creating a new client with default URL."""
+        client = Client.new_client(self.api_key)
+        self.assertIsNotNone(client)
+        self.assertEqual(client.api_key, self.api_key)
+        self.assertEqual(client.base_url, "https://api.paygent.com")
+
+    def test_new_client_with_url(self):
+        """Test creating a new client with custom URL."""
+        custom_url = "https://custom-api.paygent.com"
+        client = Client.new_client_with_url(self.api_key, custom_url)
+        self.assertIsNotNone(client)
+        self.assertEqual(client.api_key, self.api_key)
+        self.assertEqual(client.base_url, custom_url)
+
+    def test_calculate_cost_llama(self):
+        """Test cost calculation for Llama model."""
+        usage_data = UsageData(
+            service_provider="openai",
+            model="llama",
+            prompt_tokens=1000,
+            completion_tokens=500,
+            total_tokens=1500
+        )
+        
+        cost = self.client._calculate_cost("llama", usage_data)
+        expected = 0.00015  # (1000/1000 + 500/1000) * 0.0001
+        self.assertAlmostEqual(cost, expected, places=6)
+
+    def test_calculate_cost_gpt4(self):
+        """Test cost calculation for GPT-4 model."""
+        usage_data = UsageData(
+            service_provider="openai",
+            model="gpt-4",
+            prompt_tokens=1000,
+            completion_tokens=500,
+            total_tokens=1500
+        )
+        
+        cost = self.client._calculate_cost("gpt-4", usage_data)
+        expected = 0.06  # (1000/1000) * 0.03 + (500/1000) * 0.06
+        self.assertEqual(cost, expected)
+
+    def test_calculate_cost_unknown_model(self):
+        """Test cost calculation for unknown model (default pricing)."""
+        usage_data = UsageData(
+            service_provider="custom",
+            model="unknown-model",
+            prompt_tokens=1000,
+            completion_tokens=500,
+            total_tokens=1500
+        )
+        
+        cost = self.client._calculate_cost("unknown-model", usage_data)
+        expected = 0.00015  # (1000/1000 + 500/1000) * 0.0001 (default pricing)
+        self.assertAlmostEqual(cost, expected, places=6)
+
+    @patch('paygent_sdk.client.requests.Session.post')
+    def test_send_usage_success(self, mock_post):
+        """Test successful usage data sending."""
+        # Mock successful response
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.text = "Success"
+        mock_post.return_value = mock_response
+
+        usage_data = UsageData(
+            service_provider="openai",
+            model="llama",
+            prompt_tokens=1000,
+            completion_tokens=500,
+            total_tokens=1500
+        )
+
+        # Should not raise any exception
+        self.client.send_usage("agent-123", "customer-456", "test-indicator", usage_data)
+
+        # Verify the request was made correctly
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        
+        # Check URL
+        self.assertEqual(call_args[0][0], "https://api.paygent.com/api/v1/usage")
+        
+        # Check headers
+        headers = call_args[1]['headers']
+        self.assertEqual(headers['Content-Type'], 'application/json')
+        self.assertEqual(headers['paygent-api-key'], 'test-api-key')
+        
+        # Check request data
+        request_data = call_args[1]['json']
+        self.assertEqual(request_data['agentId'], 'agent-123')
+        self.assertEqual(request_data['customerId'], 'customer-456')
+        self.assertEqual(request_data['indicator'], 'test-indicator')
+        self.assertAlmostEqual(request_data['amount'], 0.00015, places=6)  # Expected cost for llama model
+        self.assertEqual(request_data['inputToken'], 1000)  # prompt_tokens
+        self.assertEqual(request_data['outputToken'], 500)  # completion_tokens
+        self.assertEqual(request_data['model'], 'llama')  # model
+        self.assertEqual(request_data['serviceProvider'], 'openai')  # service_provider
+
+    @patch('paygent_sdk.client.requests.Session.post')
+    def test_send_usage_http_error(self, mock_post):
+        """Test handling of HTTP errors."""
+        # Mock HTTP error response
+        mock_response = Mock()
+        mock_response.status_code = 400
+        mock_response.text = "Bad Request"
+        mock_response.raise_for_status.side_effect = requests.HTTPError("400 Bad Request")
+        mock_post.return_value = mock_response
+
+        usage_data = UsageData(
+            service_provider="openai",
+            model="llama",
+            prompt_tokens=1000,
+            completion_tokens=500,
+            total_tokens=1500
+        )
+
+        # Should raise HTTPError
+        with self.assertRaises(requests.HTTPError):
+            self.client.send_usage("agent-123", "customer-456", "test-indicator", usage_data)
+
+    @patch('paygent_sdk.client.requests.Session.post')
+    def test_send_usage_network_error(self, mock_post):
+        """Test handling of network errors."""
+        # Mock network error
+        mock_post.side_effect = requests.ConnectionError("Network error")
+
+        usage_data = UsageData(
+            service_provider="openai",
+            model="llama",
+            prompt_tokens=1000,
+            completion_tokens=500,
+            total_tokens=1500
+        )
+
+        # Should raise ConnectionError
+        with self.assertRaises(requests.ConnectionError):
+            self.client.send_usage("agent-123", "customer-456", "test-indicator", usage_data)
+
+    def test_set_log_level(self):
+        """Test setting log level."""
+        self.client.set_log_level(logging.DEBUG)
+        self.assertEqual(self.client.logger.level, logging.DEBUG)
+
+    def test_get_logger(self):
+        """Test getting logger instance."""
+        logger = self.client.get_logger()
+        self.assertIsNotNone(logger)
+        self.assertEqual(logger, self.client.logger)
+
+    @patch('paygent_sdk.client.requests.Session.post')
+    def test_send_usage_with_token_string_success(self, mock_post):
+        """Test successful usage data sending with token strings."""
+        # Mock successful response
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.text = "Success"
+        mock_post.return_value = mock_response
+
+        usage_data = UsageDataWithStrings(
+            service_provider="openai",
+            model="llama",
+            prompt_string="Hello, world!",
+            output_string="Hi there!"
+        )
+
+        # Should not raise any exception
+        self.client.send_usage_with_token_string("agent-123", "customer-456", "test-indicator", usage_data)
+
+        # Verify the request was made correctly
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        
+        # Check URL
+        self.assertEqual(call_args[0][0], "https://api.paygent.com/api/v1/usage")
+        
+        # Check headers
+        headers = call_args[1]['headers']
+        self.assertEqual(headers['Content-Type'], 'application/json')
+        self.assertEqual(headers['paygent-api-key'], 'test-api-key')
+        
+        # Check request data
+        request_data = call_args[1]['json']
+        self.assertEqual(request_data['agentId'], 'agent-123')
+        self.assertEqual(request_data['customerId'], 'customer-456')
+        self.assertEqual(request_data['indicator'], 'test-indicator')
+        self.assertIn('amount', request_data)  # Amount should be calculated
+        self.assertIn('inputToken', request_data)  # inputToken should be present
+        self.assertIn('outputToken', request_data)  # outputToken should be present
+        self.assertIsInstance(request_data['inputToken'], int)  # inputToken should be integer
+        self.assertIsInstance(request_data['outputToken'], int)  # outputToken should be integer
+        self.assertEqual(request_data['model'], 'llama')  # model
+        self.assertEqual(request_data['serviceProvider'], 'openai')  # service_provider
+
+
+class TestModels(unittest.TestCase):
+    """Test cases for data models."""
+
+    def test_usage_data(self):
+        """Test UsageData model."""
+        usage_data = UsageData(
+            service_provider="openai",
+            model="llama",
+            prompt_tokens=1000,
+            completion_tokens=500,
+            total_tokens=1500
+        )
+        
+        self.assertEqual(usage_data.service_provider, "openai")
+        self.assertEqual(usage_data.model, "llama")
+        self.assertEqual(usage_data.prompt_tokens, 1000)
+        self.assertEqual(usage_data.completion_tokens, 500)
+        self.assertEqual(usage_data.total_tokens, 1500)
+
+    def test_api_request(self):
+        """Test APIRequest model."""
+        api_request = APIRequest(
+            agent_id="agent-123",
+            customer_id="customer-456",
+            indicator="test-indicator",
+            amount=0.15
+        )
+        
+        self.assertEqual(api_request.agent_id, "agent-123")
+        self.assertEqual(api_request.customer_id, "customer-456")
+        self.assertEqual(api_request.indicator, "test-indicator")
+        self.assertEqual(api_request.amount, 0.15)
+
+    def test_model_pricing(self):
+        """Test ModelPricing model."""
+        pricing = ModelPricing(
+            prompt_tokens_cost=0.0001,
+            completion_tokens_cost=0.0001
+        )
+        
+        self.assertEqual(pricing.prompt_tokens_cost, 0.0001)
+        self.assertEqual(pricing.completion_tokens_cost, 0.0001)
+
+    def test_model_pricing_constants(self):
+        """Test that MODEL_PRICING contains expected models."""
+        expected_models = [
+            "llama", "gpt-3.5-turbo", "gpt-4", 
+            "claude-3-sonnet", "claude-3-opus"
+        ]
+        
+        for model in expected_models:
+            self.assertIn(model, MODEL_PRICING)
+            self.assertIsInstance(MODEL_PRICING[model], ModelPricing)
+
+
+if __name__ == '__main__':
+    unittest.main()