fix: test cases failure

vesoft-inc · Jan 8, 2024 · bfb858a · bfb858a
1 parent e13a15b
commit bfb858a
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 22 deletions.
diff --git a/README.md b/README.md
@@ -3,6 +3,9 @@
 This repository holds the official Python API for NebulaGraph.
 
 [![pdm-managed](https://img.shields.io/badge/pdm-managed-blueviolet)](https://pdm.fming.dev)
+[![PyPI](https://img.shields.io/pypi/v/nebula3-python?logo=python&logoColor=%23cccccc)](https://pypi.org/project/nebula3-python)
+[![Python Version](https://img.shields.io/pypi/pyversions/nebula3-python)](https://pypi.org/project/nebula3-python)
+
 
 ## Before you start
 

diff --git a/nebula3/gclient/net/Connection.py b/nebula3/gclient/net/Connection.py
@@ -187,12 +187,6 @@ def execute_json_with_parameter(self, session_id, stmt, params):
         """
         try:
             resp = self._connection.executeJsonWithParameter(session_id, stmt, params)
-            if resp.error_code == ErrorCode.E_SESSION_INVALID:
-                raise SessionException(resp.error_code, resp.error_msg)
-            if resp.error_code == ErrorCode.E_SESSION_TIMEOUT:
-                raise SessionException(resp.error_code, resp.error_msg)
-            if resp.error_code == ErrorCode.E_EXECUTION_ERROR:
-                raise ExecutionErrorException(resp.error_msg)
             return resp
         except Exception as te:
             if isinstance(te, TTransportException):

diff --git a/nebula3/gclient/net/SessionPool.py b/nebula3/gclient/net/SessionPool.py
@@ -180,7 +180,7 @@ def execute_parameter(self, stmt, params):
                 session = self._get_idle_session()
                 if session is None:
                     raise RuntimeError('Get session failed')
-                self._add_session_to_active(session)
+                self._add_session_to_idle(session)
             raise se
 
         except Exception as e:
@@ -279,7 +279,7 @@ def execute_json_with_parameter(self, stmt, params):
                 session = self._get_idle_session()
                 if session is None:
                     raise RuntimeError('Get session failed')
-                self._add_session_to_active(session)
+                self._add_session_to_idle(session)
             raise se
         except Exception as e:
             logger.error('Execute failed: {}'.format(e))

diff --git a/tests/test_session.py b/tests/test_session.py
@@ -99,33 +99,32 @@ def test_5_session_exception(self):
         try:
             session = self.pool.get_session(self.user_name, self.password)
             another_session = self.pool.get_session(self.user_name, self.password)
-            session_id = session.session_id
+            session_id = session._session_id
             another_session.execute(f"KILL SESSION {session_id}")
             session.execute("SHOW HOSTS")
         except Exception as ex:
             assert isinstance(ex, SessionException), "expect to get SessionException"
 
     def test_6_execute_exception(self):
         # test ExecutionErrorException will be raised when execute error
-        # we need to mock a query's response code to -1005
-        import unittest
-        from unittest.mock import call
+        # we need to mock a query's response code to trigger ExecutionErrorException
+        from unittest.mock import Mock, patch
 
         try:
             session = self.pool.get_session(self.user_name, self.password)
-            # Mocking a remote call that will trigger an ExecutionErrorException
-            with unittest.mock.patch(
-                "nebula3.gclient.net.Connection._connection.executeWithParameter"
+            # Mocking the Connection.execute_parameter method
+            with patch(
+                'nebula3.graph.GraphService.Client.executeWithParameter'
             ) as mock_execute:
-                mock_execute.return_value.error_code = (
-                    ExecutionErrorException.E_EXECUTION_ERROR
-                )
+                mock_response = Mock()
+                mock_response.error_code = ExecutionErrorException.E_EXECUTION_ERROR
+                mock_execute.return_value = mock_response
                 session.execute("SHOW HOSTS")
-                # Assert that executeWithParameter was called 3 times (retry mechanism)
+                # Assert that execute_parameter was called 3 times (retry mechanism)
                 assert (
                     mock_execute.call_count == 3
-                ), "executeWithParameter was not retried 3 times"
+                ), "execute_parameter was not retried 3 times"
         except ExecutionErrorException as ex:
             assert True, "ExecutionErrorException triggered as expected"
         except Exception as ex:
-            assert False, "Unexpected exception: " + str(ex)
+            assert False, f"Unexpected exception: {str(ex)}"
diff --git a/tests/test_session_pool.py b/tests/test_session_pool.py
@@ -156,7 +156,7 @@ def test_session_renew_when_invalid(self):
 
         # kill all sessions of the pool, size 1 here though
         for session in session_pool._idle_sessions:
-            session_id = session.session_id
+            session_id = session._session_id
             session.execute(f"KILL SESSION {session_id}")
         try:
             session_pool.execute("SHOW HOSTS;")
@@ -170,6 +170,7 @@ def test_session_renew_when_invalid(self):
         ), "session should be renewed"
         resp = session_pool.execute("SHOW HOSTS;")
         assert resp.is_succeeded(), "session_pool should be usable after renewing"
+        session_pool.close()
 
 
 def test_session_pool_multi_thread():