"""Unit tests for neo4j graph."""

from unittest.mock import ANY, Mock

from src.model.neo4j_graph import neo4j_session, run


class TestNeo4jGraph:
    """Tests for neo4j graph functions."""

    def test_neo4j_session_decorator(self, neo4j_session_mock):
        """Test get_product function."""

        @neo4j_session
        def function_example(param, neo4j_session):
            """Product id."""
            return neo4j_session.write_transaction(ANY, param)

        result_mock = neo4j_session_mock.write_transaction.return_value
        param_example = 'Cypher Code'
        result = function_example(param_example)
        assert result == result_mock
        neo4j_session_mock.write_transaction.assert_called_once_with(ANY, param_example)

    def test_run(self, neo4j_session_mock):
        """Test run function."""
        query = 'MATCH(i:Itdentity) RETURN i'
        params = {}
        expected = Mock()
        expected.data.return_value = {'foo': 'bar'}
        neo4j_session_mock.write_transaction.return_value = expected

        result_mock = run(query, params)
        assert expected == result_mock