import unittest
from unittest.mock import Mock, patch
import os
from .queryclient import QueryClient


class TestQueryClient(unittest.TestCase):
    
    def test_default_max_inbound_message_size(self):
        """Test that the default max inbound message size is set correctly"""
        expected_size = 1024 * 1024 * 64  # 64MiB
        self.assertEqual(QueryClient.DEFAULT_MAX_INBOUND_IN_BYTES, expected_size)
    
    def test_init(self):
        """Test QueryClient initialization"""
        mock_channel = Mock()
        client = QueryClient(mock_channel)
        
        self.assertEqual(client.channel, mock_channel)
        self.assertIsNotNone(client.inner)
    
    @patch.dict(os.environ, {'CONNECTOR_HOST': 'test-host'})
    def test_from_connector_id_with_env_var(self):
        """Test from_connector_id with CONNECTOR_HOST environment variable"""
        with patch.object(QueryClient, 'from_target') as mock_from_target:
            mock_client = Mock()
            mock_from_target.return_value = mock_client
            
            result = QueryClient.from_connector_id(123)
            
            expected_target = "test-host:10123"
            mock_from_target.assert_called_once_with(
                expected_target, QueryClient.DEFAULT_MAX_INBOUND_IN_BYTES
            )
            self.assertEqual(result, mock_client)
    
    @patch.dict(os.environ, {}, clear=True)
    def test_from_connector_id_without_env_var(self):
        """Test from_connector_id without CONNECTOR_HOST environment variable"""
        with patch.object(QueryClient, 'from_target') as mock_from_target:
            mock_client = Mock()
            mock_from_target.return_value = mock_client
            
            result = QueryClient.from_connector_id(456)
            
            expected_target = "localhost:10456"
            mock_from_target.assert_called_once_with(
                expected_target, QueryClient.DEFAULT_MAX_INBOUND_IN_BYTES
            )
            self.assertEqual(result, mock_client)
    
    def test_from_connector_id_with_custom_message_size(self):
        """Test from_connector_id with custom max_inbound_message_size"""
        custom_size = 1024 * 1024 * 128  # 128MiB
        
        with patch.object(QueryClient, 'from_target') as mock_from_target:
            mock_client = Mock()
            mock_from_target.return_value = mock_client
            
            result = QueryClient.from_connector_id(789, custom_size)
            
            expected_target = "localhost:10789"
            mock_from_target.assert_called_once_with(expected_target, custom_size)
            self.assertEqual(result, mock_client)
    
    def test_from_target(self):
        """Test from_target class method"""
        with patch.object(QueryClient, 'resolve') as mock_resolve:
            mock_channel = Mock()
            mock_resolve.return_value = mock_channel
            
            result = QueryClient.from_target("test-target:8080", 1024)
            
            mock_resolve.assert_called_once_with(
                target="test-target:8080", max_inbound_message_size=1024
            )
            self.assertIsInstance(result, QueryClient)
            self.assertEqual(result.channel, mock_channel)
    
    
    @patch.object(QueryClient, 'configuration_from_b64_string')
    def test_new_session_from_b64_config_only(self, mock_config_from_b64):
        """Test new_session_from_b64 with configuration only"""
        mock_channel = Mock()
        mock_config = Mock()
        
        client = QueryClient(mock_channel)
        
        # Mock configuration_from_b64_string to return a mock config
        mock_config_from_b64.return_value = mock_config
        
        # Mock new_session
        mock_session = Mock()
        client.new_session = Mock(return_value=mock_session)
        
        # Use any base64 string (we're mocking the decode anyway)
        result = client.new_session_from_b64("encoded_config")
        
        # Check that configuration_from_b64_string was called
        mock_config_from_b64.assert_called_once_with("encoded_config")
        
        # Verify new_session was called with config and None auth
        client.new_session.assert_called_once_with(mock_config, None)
        self.assertEqual(result, mock_session)
    
    @patch.object(QueryClient, 'auth_configuration_from_b64_string')
    @patch.object(QueryClient, 'configuration_from_b64_string')
    def test_new_session_from_b64_with_auth(self, mock_config_from_b64, mock_auth_from_b64):
        """Test new_session_from_b64 with configuration and authentication"""
        mock_channel = Mock()
        mock_config = Mock()
        mock_auth_obj = Mock()
        
        client = QueryClient(mock_channel)
        
        # Mock the static methods
        mock_config_from_b64.return_value = mock_config
        mock_auth_from_b64.return_value = mock_auth_obj
        
        # Mock new_session
        mock_session = Mock()
        client.new_session = Mock(return_value=mock_session)
        
        result = client.new_session_from_b64("encoded_config", "encoded_auth")
        
        # Check that both static methods were called
        mock_config_from_b64.assert_called_once_with("encoded_config")
        mock_auth_from_b64.assert_called_once_with("encoded_auth")
        
        # Verify new_session was called with both config and auth
        client.new_session.assert_called_once_with(mock_config, mock_auth_obj)
        self.assertEqual(result, mock_session)


if __name__ == '__main__':
    unittest.main()