fix Neo4jHook to get the query response

airflow/providers/neo4j/hooks/neo4j.py

@@ -52,15 +52,15 @@ def get_conn(self) -> Neo4jDriver:
         Function that initiates a new Neo4j connection
         with username, password and database schema.
         """
+        if self.client is not None:
+            return self.client
+
         self.connection = self.get_connection(self.neo4j_conn_id)
         self.extras = self.connection.extra_dejson.copy()
 
         self.uri = self.get_uri(self.connection)
         self.log.info('URI: %s', self.uri)
 
-        if self.client is not None:
-            return self.client
-
         is_encrypted = self.connection.extra_dejson.get('encrypted', False)
 
         self.client = GraphDatabase.driver(
@@ -76,6 +76,7 @@ def get_uri(self, conn: Connection) -> str:
         - neo4j_scheme - neo4j://
         - certs_self_signed - neo4j+ssc://
         - certs_trusted_ca - neo4j+s://
+
         :param conn: connection object.
         :return: uri
         """
@@ -106,15 +107,15 @@ def run(self, query) -> Result:
         Function to create a neo4j session
         and execute the query in the session.
 
-
         :param query: Neo4j query
         :return: Result
         """
         driver = self.get_conn()
         if not self.connection.schema:
             with driver.session() as session:
                 result = session.run(query)
+                return result.data()
         else:
             with driver.session(database=self.connection.schema) as session:
                 result = session.run(query)
-        return result
+                return result.data()

tests/providers/neo4j/hooks/test_neo4j.py

@@ -15,51 +15,94 @@
 # specific language governing permissions and limitations
 # under the License.
 #
-import json
 import unittest
 from unittest import mock
 
+from parameterized import parameterized
+
 from airflow.models import Connection
 from airflow.providers.neo4j.hooks.neo4j import Neo4jHook
 
 
 class TestNeo4jHookConn(unittest.TestCase):
-    def setUp(self):
-        super().setUp()
-        self.neo4j_hook = Neo4jHook()
-        self.connection = Connection(
-            conn_type='neo4j', login='login', password='password', host='host', schema='schema'
+    @parameterized.expand(
+        [
+            [{}, "bolt://host:7687"],
+            [{"bolt_scheme": True}, "bolt://host:7687"],
+            [{"certs_self_signed": True, "bolt_scheme": True}, "bolt+ssc://host:7687"],
+            [{"certs_trusted_ca": True, "bolt_scheme": True}, "bolt+s://host:7687"],
+        ]
+    )
+    def test_get_uri_neo4j_scheme(self, conn_extra, expected_uri):
+        connection = Connection(
+            conn_type='neo4j',
+            login='login',
+            password='password',
+            host='host',
+            schema='schema',
+            extra=conn_extra,
         )
 
-    def test_get_uri_neo4j_scheme(self):
-
-        self.neo4j_hook.get_connection = mock.Mock()
-        self.neo4j_hook.get_connection.return_value = self.connection
-        uri = self.neo4j_hook.get_uri(self.connection)
-
-        assert uri == "bolt://host:7687"
-
-    def test_get_uri_bolt_scheme(self):
+        # Use the environment variable mocking to test saving the configuration as a URI and
+        # to avoid mocking Airflow models class
+        with mock.patch.dict('os.environ', AIRFLOW_CONN_NEO4J_DEFAULT=connection.get_uri()):
+            neo4j_hook = Neo4jHook()
+            uri = neo4j_hook.get_uri(connection)
 
-        self.connection.extra = json.dumps({"bolt_scheme": True})
-        self.neo4j_hook.get_connection = mock.Mock()
-        self.neo4j_hook.get_connection.return_value = self.connection
-        uri = self.neo4j_hook.get_uri(self.connection)
+            assert uri == expected_uri
 
-        assert uri == "bolt://host:7687"
-
-    def test_get_uri_bolt_ssc_scheme(self):
-        self.connection.extra = json.dumps({"certs_self_signed": True, "bolt_scheme": True})
-        self.neo4j_hook.get_connection = mock.Mock()
-        self.neo4j_hook.get_connection.return_value = self.connection
-        uri = self.neo4j_hook.get_uri(self.connection)
+    @mock.patch('airflow.providers.neo4j.hooks.neo4j.GraphDatabase')
+    def test_run_with_schema(self, mock_graph_database):
+        connection = Connection(
+            conn_type='neo4j', login='login', password='password', host='host', schema='schema'
+        )
+        mock_sql = mock.MagicMock(name="sql")
 
-        assert uri == "bolt+ssc://host:7687"
+        # Use the environment variable mocking to test saving the configuration as a URI and
+        # to avoid mocking Airflow models class
+        with mock.patch.dict('os.environ', AIRFLOW_CONN_NEO4J_DEFAULT=connection.get_uri()):
+            neo4j_hook = Neo4jHook()
+            op_result = neo4j_hook.run(mock_sql)
+            mock_graph_database.assert_has_calls(
+                [
+                    mock.call.driver('bolt://host:7687', auth=('login', 'password'), encrypted=False),
+                    mock.call.driver().session(database='schema'),
+                    mock.call.driver().session().__enter__(),
+                    mock.call.driver().session().__enter__().run(mock_sql),
+                    mock.call.driver().session().__enter__().run().data(),
+                    mock.call.driver().session().__exit__(None, None, None),
+                ]
+            )
+            session = mock_graph_database.driver.return_value.session.return_value.__enter__.return_value
+            self.assertEqual(
+                session.run.return_value.data.return_value,
+                op_result,
+            )
 
-    def test_get_uri_bolt_trusted_ca_scheme(self):
-        self.connection.extra = json.dumps({"certs_trusted_ca": True, "bolt_scheme": True})
-        self.neo4j_hook.get_connection = mock.Mock()
-        self.neo4j_hook.get_connection.return_value = self.connection
-        uri = self.neo4j_hook.get_uri(self.connection)
+    @mock.patch('airflow.providers.neo4j.hooks.neo4j.GraphDatabase')
+    def test_run_without_schema(self, mock_graph_database):
+        connection = Connection(
+            conn_type='neo4j', login='login', password='password', host='host', schema=None
+        )
+        mock_sql = mock.MagicMock(name="sql")
 
-        assert uri == "bolt+s://host:7687"
+        # Use the environment variable mocking to test saving the configuration as a URI and
+        # to avoid mocking Airflow models class
+        with mock.patch.dict('os.environ', AIRFLOW_CONN_NEO4J_DEFAULT=connection.get_uri()):
+            neo4j_hook = Neo4jHook()
+            op_result = neo4j_hook.run(mock_sql)
+            mock_graph_database.assert_has_calls(
+                [
+                    mock.call.driver('bolt://host:7687', auth=('login', 'password'), encrypted=False),
+                    mock.call.driver().session(),
+                    mock.call.driver().session().__enter__(),
+                    mock.call.driver().session().__enter__().run(mock_sql),
+                    mock.call.driver().session().__enter__().run().data(),
+                    mock.call.driver().session().__exit__(None, None, None),
+                ]
+            )
+            session = mock_graph_database.driver.return_value.session.return_value.__enter__.return_value
+            self.assertEqual(
+                session.run.return_value.data.return_value,
+                op_result,
+            )

tests/providers/neo4j/operators/test_neo4j.py

@@ -15,27 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
-# specific language governing permissions and limitations
-# under the License.
 import unittest
 from unittest import mock
 
-from airflow.models.dag import DAG
 from airflow.providers.neo4j.operators.neo4j import Neo4jOperator
 from airflow.utils import timezone
 
@@ -46,16 +28,13 @@
 
 
 class TestNeo4jOperator(unittest.TestCase):
-    def setUp(self):
-        args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
-        dag = DAG(TEST_DAG_ID, default_args=args)
-        self.dag = dag
-
-    @mock.patch('airflow.providers.neo4j.operators.neo4j.Neo4jOperator.get_hook')
+    @mock.patch('airflow.providers.neo4j.operators.neo4j.Neo4jHook')
     def test_neo4j_operator_test(self, mock_hook):
 
         sql = """
             MATCH (tom {name: "Tom Hanks"}) RETURN tom
             """
-        op = Neo4jOperator(task_id='basic_neo4j', sql=sql, dag=self.dag)
-        op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
+        op = Neo4jOperator(task_id='basic_neo4j', sql=sql)
+        op.execute(mock.MagicMock())
+        mock_hook.assert_called_once_with(conn_id='neo4j_default')
+        mock_hook.return_value.run.assert_called_once_with(sql)