"""Tests for Snowflake connector wrapper with session token support."""

import pytest
from unittest.mock import Mock, patch, MagicMock
from test_runner.validate.connector_wrapper import (
    SnowflakeConnectorWrapper,
    create_connector_from_connection
)


class TestSnowflakeConnectorWrapper:
    """Test suite for SnowflakeConnectorWrapper."""

    def test_init_with_default_ownership(self):
        """Test initialization with default ownership (does not own connection)."""
        mock_connection = Mock()
        wrapper = SnowflakeConnectorWrapper(mock_connection)
        
        assert wrapper._connection == mock_connection
        assert wrapper._owns_connection is False
        assert wrapper._closed is False

    def test_init_with_explicit_ownership(self):
        """Test initialization with explicit ownership flag."""
        mock_connection = Mock()
        wrapper = SnowflakeConnectorWrapper(mock_connection, owns_connection=True)
        
        assert wrapper._connection == mock_connection
        assert wrapper._owns_connection is True
        assert wrapper._closed is False

    def test_connection_property_returns_connection(self):
        """Test that connection property returns wrapped connection."""
        mock_connection = Mock()
        wrapper = SnowflakeConnectorWrapper(mock_connection)
        
        assert wrapper.connection == mock_connection

    def test_connection_property_raises_when_closed(self):
        """Test that accessing connection after close raises RuntimeError."""
        mock_connection = Mock()
        wrapper = SnowflakeConnectorWrapper(mock_connection)
        wrapper._closed = True
        
        with pytest.raises(RuntimeError, match="Connection wrapper has been closed"):
            _ = wrapper.connection

    def test_execute_query_returns_results(self):
        """Test execute_query delegates to connection and returns results."""
        mock_connection = Mock()
        mock_cursor = Mock()
        mock_connection.cursor.return_value = mock_cursor
        
        expected_results = [{"col1": "value1"}, {"col1": "value2"}]
        mock_cursor.fetchall.return_value = expected_results
        
        wrapper = SnowflakeConnectorWrapper(mock_connection)
        results = wrapper.execute_query("SELECT * FROM table")
        
        assert results == expected_results
        mock_cursor.execute.assert_called_once_with("SELECT * FROM table")
        mock_cursor.close.assert_called_once()

    def test_execute_query_closes_cursor_on_exception(self):
        """Test that cursor is closed even if execute raises exception."""
        mock_connection = Mock()
        mock_cursor = Mock()
        mock_connection.cursor.return_value = mock_cursor
        mock_cursor.execute.side_effect = Exception("Query failed")
        
        wrapper = SnowflakeConnectorWrapper(mock_connection)
        
        with pytest.raises(Exception, match="Query failed"):
            wrapper.execute_query("SELECT * FROM table")
        
        mock_cursor.close.assert_called_once()

    def test_execute_statement_executes_ddl(self):
        """Test execute_statement delegates to connection without returning results."""
        mock_connection = Mock()
        mock_cursor = Mock()
        mock_connection.cursor.return_value = mock_cursor
        
        wrapper = SnowflakeConnectorWrapper(mock_connection)
        wrapper.execute_statement("CREATE TABLE test (id INT)")
        
        mock_cursor.execute.assert_called_once_with("CREATE TABLE test (id INT)")
        mock_cursor.close.assert_called_once()

    def test_execute_statement_closes_cursor_on_exception(self):
        """Test that cursor is closed even if execute raises exception."""
        mock_connection = Mock()
        mock_cursor = Mock()
        mock_connection.cursor.return_value = mock_cursor
        mock_cursor.execute.side_effect = Exception("DDL failed")
        
        wrapper = SnowflakeConnectorWrapper(mock_connection)
        
        with pytest.raises(Exception, match="DDL failed"):
            wrapper.execute_statement("CREATE TABLE test (id INT)")
        
        mock_cursor.close.assert_called_once()

    def test_close_without_ownership_does_not_close_connection(self):
        """Test close() when wrapper doesn't own connection."""
        mock_connection = Mock()
        wrapper = SnowflakeConnectorWrapper(mock_connection, owns_connection=False)
        
        wrapper.close()
        
        assert wrapper._closed is True
        mock_connection.close.assert_not_called()

    def test_close_with_ownership_closes_connection(self):
        """Test close() when wrapper owns connection."""
        mock_connection = Mock()
        wrapper = SnowflakeConnectorWrapper(mock_connection, owns_connection=True)
        
        wrapper.close()
        
        assert wrapper._closed is True
        mock_connection.close.assert_called_once()

    def test_close_handles_connection_close_exception(self):
        """Test close() handles exception when closing owned connection."""
        mock_connection = Mock()
        mock_connection.close.side_effect = Exception("Close failed")
        wrapper = SnowflakeConnectorWrapper(mock_connection, owns_connection=True)
        
        # Should not raise exception
        wrapper.close()
        
        assert wrapper._closed is True

    def test_close_idempotent(self):
        """Test that calling close() multiple times is safe."""
        mock_connection = Mock()
        wrapper = SnowflakeConnectorWrapper(mock_connection, owns_connection=True)
        
        wrapper.close()
        wrapper.close()  # Should not raise or double-close
        
        assert wrapper._closed is True
        assert mock_connection.close.call_count == 1

    def test_context_manager_enter(self):
        """Test context manager __enter__ returns self."""
        mock_connection = Mock()
        wrapper = SnowflakeConnectorWrapper(mock_connection)
        
        with wrapper as ctx:
            assert ctx is wrapper

    def test_context_manager_exit_calls_close(self):
        """Test context manager __exit__ calls close."""
        mock_connection = Mock()
        wrapper = SnowflakeConnectorWrapper(mock_connection, owns_connection=True)
        
        with wrapper:
            pass
        
        assert wrapper._closed is True
        mock_connection.close.assert_called_once()

    def test_create_connector_from_connection_default_ownership(self):
        """Test factory function with default ownership."""
        mock_connection = Mock()
        connector = create_connector_from_connection(mock_connection)
        
        assert isinstance(connector, SnowflakeConnectorWrapper)
        assert connector._connection == mock_connection
        assert connector._owns_connection is False

    def test_create_connector_from_connection_with_ownership(self):
        """Test factory function with explicit ownership."""
        mock_connection = Mock()
        connector = create_connector_from_connection(mock_connection, owns_connection=True)
        
        assert isinstance(connector, SnowflakeConnectorWrapper)
        assert connector._connection == mock_connection
        assert connector._owns_connection is True
