"""Test for neo4j DB session."""

import importlib
from ssl import SSLError
from unittest import mock

import connector_neo4j

from flask import Flask

from neo4j.exceptions import ServiceUnavailable
from neo4j.exceptions import SessionExpired

import pytest

app = Flask(__name__)


@mock.patch("neo4j.GraphDatabase.driver")
def test_read_connection(driver_create):
    """Test standard read only session.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    with app.test_request_context() as context:
        session_create = driver_create.return_value.session
        session = session_create.return_value
        session._transaction = False
        assert not session_create.called

        with connector_neo4j.Neo4jSession():
            assert session_create.called
            assert session_create.call_count == 1
            assert session_create.call_args_list == [
                mock.call(default_access_mode="READ", database=None)
            ]
            assert not connector_neo4j.neo4j_session
            assert context.g.neo4j_session == session
            assert connector_neo4j.get_session() == session
            assert not session.close.called
            assert not session.begin_transaction.called
        assert session.close.called
        assert not session.commit.called
        assert not session.rollback.called
        assert context.g.neo4j_session is None


@mock.patch("neo4j.GraphDatabase.driver")
def test_force_write_connection(driver_create):
    """Test override non-transaction session to write server.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    with app.test_request_context() as context:
        session_create = driver_create.return_value.session
        session = session_create.return_value
        session._transaction = False
        assert not session_create.called

        with connector_neo4j.Neo4jSession(force_write_server=True):
            assert session_create.called
            assert session_create.call_count == 1
            assert session_create.call_args_list == [
                mock.call(default_access_mode="WRITE", database=None)
            ]
            assert not connector_neo4j.neo4j_session
            assert context.g.neo4j_session == session
            assert connector_neo4j.get_session() == session
            assert not session.close.called
            assert not session.begin_transaction.called
        assert session.close.called
        assert not session.commit.called
        assert not session.rollback.called
        assert context.g.neo4j_session is None


@mock.patch("neo4j.GraphDatabase.driver")
def test_commit(driver_create):
    """Test open and commit transaction.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    with app.test_request_context() as context:
        session_create = driver_create.return_value.session
        session = session_create.return_value
        transaction = session._transaction

        assert "neo4j_session" not in context.g
        with connector_neo4j.Neo4jSession(transaction=True):
            assert session_create.called
            assert session_create.call_count == 1
            assert session_create.call_args_list == [
                mock.call(default_access_mode="WRITE", database=None)
            ]
            assert context.g.neo4j_session == session
            assert connector_neo4j.get_session() == transaction
            assert session.begin_transaction.called
            assert not session.close.called
            assert not transaction.commit.called
        assert session.close.called
        assert transaction.commit.called
        assert not transaction.rollback.called

        # Verify get_session handles context.g.neo4j_session is None
        assert context.g.neo4j_session is None
        with connector_neo4j.Neo4jSession(transaction=True):
            pass


@mock.patch("neo4j.GraphDatabase.driver")
def test_rollback(driver_create):
    """Test open and rollback transaction.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    with app.test_request_context() as context:
        session_create = driver_create.return_value.session
        session = session_create.return_value
        transaction = session._transaction

        @connector_neo4j.Neo4jSession(transaction=True)
        def func():
            assert session_create.called
            assert session_create.call_count == 1
            assert session_create.call_args_list == [
                mock.call(default_access_mode="WRITE", database=None)
            ]
            assert context.g.neo4j_session == transaction
            assert connector_neo4j.get_session() == transaction
            assert session.begin_transaction.called
            assert not session.close.called
            assert not transaction.rollback.called
            raise Exception("something went wrong")

        try:
            func()
        except:  # noqa: E722
            pass
        assert session.close.called
        assert not transaction.commit.called

        # Verify get_session handles context.g.neo4j_session is None
        assert context.g.neo4j_session is None
        with connector_neo4j.Neo4jSession(transaction=True):
            pass


@mock.patch("neo4j.GraphDatabase.driver")
def test_commit_failure(driver_create):
    """Test error during transaction commit.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    commit_error_msg = "commit error"
    with app.test_request_context() as context:
        session_create = driver_create.return_value.session
        session = session_create.return_value
        transaction = session._transaction
        transaction.commit.side_effect = Exception(commit_error_msg)

        @connector_neo4j.Neo4jSession(transaction=True)
        def func():
            pass

        with pytest.raises(Exception) as exc:
            func()
        assert session.close.called
        assert str(exc.value) == commit_error_msg
        assert context.g.neo4j_session is None


@mock.patch("neo4j.GraphDatabase.driver")
def test_create_session_twice(driver_create):
    """Test creating session inside of session.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    with app.test_request_context():

        @connector_neo4j.Neo4jSession()
        def func_one():
            pass

        @connector_neo4j.Neo4jSession()
        def func_two():
            func_one()

        with pytest.raises(connector_neo4j.exceptions.SessionAlreadyOpen):
            func_two()


@mock.patch("neo4j.GraphDatabase.driver")
def test_default_retry(driver_create):
    """Test exhausting pool according to size."""
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    session_create = driver_create.return_value.session
    session = session_create.return_value

    pool_start_size = 10
    session._pool.address = "localhost"
    session._pool.connections = {
        "localhost": [mock.MagicMock() for x in range(0, pool_start_size)]
    }

    # simular connection failed and removal from pool
    def mock_session_failed(arg):
        connections = session._pool.connections["localhost"]
        if connections:
            connections.pop()
        raise SessionExpired("test")

    session.run.side_effect = mock_session_failed

    with app.test_request_context():

        @connector_neo4j.Neo4jSession()
        def func():
            pass

        with pytest.raises(SessionExpired):
            func()
        assert session.run.call_count == pool_start_size + 1


@mock.patch("neo4j.GraphDatabase.driver")
def test_auto_refresh_session(driver_create):
    """Test session retry on expiration.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    session_create = driver_create.return_value.session
    session = session_create.return_value
    session._transaction = False
    session.run.side_effect = [
        SessionExpired("test"),
        ServiceUnavailable,
        mock.MagicMock(),
    ]

    with app.test_request_context() as context:

        @connector_neo4j.Neo4jSession()
        def func():
            assert session_create.call_count == 3
            assert context.g.neo4j_session == session
            assert connector_neo4j.get_session() == session

        func()


@pytest.mark.parametrize(
    "root_error, final_error_type",
    [
        (OSError(9, "Bad file descriptor"), ServiceUnavailable),
        (SSLError(), ServiceUnavailable),
        (
            AttributeError("'NoneType' object has no attribute 'complete'"),
            ServiceUnavailable,
        ),  # noqa:E501
        (OSError(9, "Blerg"), OSError),
        (OSError(9), OSError),
    ],
)
@mock.patch("neo4j.GraphDatabase.driver")
def test_session_create_mapped_exceptions(driver_create, root_error, final_error_type):
    """Test session failure on create special exceptions.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    session_create = driver_create.return_value.session
    session = session_create.return_value
    session.run.side_effect = root_error
    session._pool.connections = {}

    with app.test_request_context():

        @connector_neo4j.Neo4jSession()
        def func():
            pass

        with pytest.raises(final_error_type):
            func()


@pytest.mark.parametrize(
    "root_error, final_error_type",
    [
        (OSError(9, "Bad file descriptor"), ServiceUnavailable),
        (SSLError(), ServiceUnavailable),
        (OSError(9, "Blerg"), OSError),
        (OSError(9), OSError),
    ],
)
@mock.patch("neo4j.GraphDatabase.driver")
def test_session_run_mapped_exceptions(driver_create, root_error, final_error_type):
    """Test exception handling during context open.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    session_create = driver_create.return_value.session
    session = session_create.return_value
    session._pool.connections = {}

    with app.test_request_context():

        @connector_neo4j.Neo4jSession()
        def func():
            raise root_error

        with pytest.raises(final_error_type):
            func()


@pytest.mark.parametrize(
    "root_error, exc_line, exc_file, final_error_type",
    [
        (
            TypeError,
            980,
            "/var/venv/lib64/python3.8/site-packages/neo4j/io/__init__.py",
            ServiceUnavailable,
        ),
        (
            TypeError,
            980,
            "/var/venv/lib32/python3.8/site-packages/neo4j/io/__init__.py",
            ServiceUnavailable,
        ),
        (
            AttributeError,
            980,
            "/var/venv/lib64/python3.8/site-packages/neo4j/io/__init__.py",
            AttributeError,
        ),
        (TypeError, 980, "__init__.py", TypeError),
        (
            TypeError,
            777,
            "/var/venv/lib64/python3.8/site-packages/neo4j/io/__init__.py",
            TypeError,
        ),
        (
            TypeError,
            980,
            "/var/venv/lib64/python3.8/site-packages/sql/io/__init__.py",
            TypeError,
        ),
    ],
)
@mock.patch("sys.exc_info")
@mock.patch("neo4j.GraphDatabase.driver")
def test_session_run_traceback_exception(
    driver_create, mock_exc_info, root_error, exc_line, exc_file, final_error_type
):
    """Test exception handling during query run using traceback.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    session_create = driver_create.return_value.session
    session = session_create.return_value
    session._pool.connections = {}

    mock_file_info = mock.MagicMock()
    mock_file_info.tb_lineno = exc_line
    mock_file_info.tb_frame.f_code.co_filename = exc_file
    mock_exc_info.return_value = (None, None, mock_file_info)

    with app.test_request_context():

        @connector_neo4j.Neo4jSession()
        def func():
            raise root_error

        with pytest.raises(final_error_type):
            func()


@mock.patch("neo4j.GraphDatabase.driver")
def test_session_create_fail(driver_create):
    """Test session failure on create.

    Args:
        driver_create (MagicMock): mock of neo4j driver constructor
    """
    importlib.reload(connector_neo4j)
    num_retries = 5
    connector_neo4j.configure(
        "url", "username", "password", connector_neo4j.SessionStorageMode.FLASK
    )
    session_create = driver_create.return_value.session
    session = session_create.return_value
    session.run.side_effect = [
        SessionExpired("test") for _ in range(0, num_retries + 1)
    ]
    session._pool.connections = {}

    with app.test_request_context():

        @connector_neo4j.Neo4jSession()
        def func():
            pass

        with pytest.raises(SessionExpired):
            func()
        assert session.run.call_count == num_retries + 1


def test_missing_session():
    """Test access session without decorator."""
    with app.test_request_context():
        with pytest.raises(connector_neo4j.exceptions.SessionNotCreated):
            connector_neo4j.get_session()