test: use sqlite3 session in test_api (#38736)

api/tests/unit_tests/core/external_data_tool/api/test_api.py

@@ -1,22 +1,51 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import pytest
+from sqlalchemy.orm import Session
 
+import core.external_data_tool.api.api as api_module
 from core.external_data_tool.api.api import ApiExternalDataTool
-from models.api_based_extension import APIBasedExtensionPoint
+from models.api_based_extension import APIBasedExtension, APIBasedExtensionPoint
+
+pytestmark = [
+    pytest.mark.usefixtures("sqlite_session"),
+    pytest.mark.parametrize("sqlite_session", [(APIBasedExtension,)], indirect=True),
+]
+
+
+class _DatabaseBinding:
+    """Expose the real SQLite session used by extension queries."""
+
+    session: Session
+
+    def __init__(self, session: Session) -> None:
+        self.session = session
+
+
+@pytest.fixture(autouse=True)
+def bind_sqlite_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(api_module, "db", _DatabaseBinding(sqlite_session))
+
+
+@pytest.fixture
+def api_extension(sqlite_session: Session) -> APIBasedExtension:
+    extension = APIBasedExtension(
+        tenant_id="tenant_id",
+        name="Test extension",
+        api_endpoint="http://api",
+        api_key="encrypted_key",
+    )
+    extension.id = "ext_id"
+    sqlite_session.add(extension)
+    sqlite_session.commit()
+    return extension
 
 
 def test_api_external_data_tool_name():
     assert ApiExternalDataTool.name == "api"
 
 
-@patch("core.external_data_tool.api.api.db")
-def test_validate_config_success(mock_db):
-    mock_extension = MagicMock()
-    mock_extension.id = "ext_id"
-    mock_extension.tenant_id = "tenant_id"
-    mock_db.session.scalar.return_value = mock_extension
-
+def test_validate_config_success(api_extension: APIBasedExtension):
     # Should not raise exception
     ApiExternalDataTool.validate_config("tenant_id", {"api_based_extension_id": "ext_id"})
 
@@ -26,10 +55,7 @@ def test_validate_config_missing_id():
         ApiExternalDataTool.validate_config("tenant_id", {})
 
 
-@patch("core.external_data_tool.api.api.db")
-def test_validate_config_invalid_id(mock_db):
-    mock_db.session.scalar.return_value = None
-
+def test_validate_config_invalid_id():
     with pytest.raises(ValueError, match="api_based_extension_id is invalid"):
         ApiExternalDataTool.validate_config("tenant_id", {"api_based_extension_id": "ext_id"})
 
@@ -42,16 +68,9 @@ def api_tool():
     )
 
 
-@patch("core.external_data_tool.api.api.db")
 @patch("core.external_data_tool.api.api.encrypter")
 @patch("core.external_data_tool.api.api.APIBasedExtensionRequestor")
-def test_query_success(mock_requestor_class, mock_encrypter, mock_db, api_tool):
-    mock_extension = MagicMock()
-    mock_extension.id = "ext_id"
-    mock_extension.tenant_id = "tenant_id"
-    mock_extension.api_endpoint = "http://api"
-    mock_extension.api_key = "encrypted_key"
-    mock_db.session.scalar.return_value = mock_extension
+def test_query_success(mock_requestor_class, mock_encrypter, api_tool, api_extension: APIBasedExtension):
     mock_encrypter.decrypt_token.return_value = "decrypted_key"
 
     mock_requestor = mock_requestor_class.return_value
@@ -81,24 +100,14 @@ def test_query_missing_extension_id():
         api_tool.query({}, "")
 
 
-@patch("core.external_data_tool.api.api.db")
-def test_query_invalid_extension(mock_db, api_tool):
-    mock_db.session.scalar.return_value = None
-
+def test_query_invalid_extension(api_tool):
     with pytest.raises(ValueError, match=".*error: api_based_extension_id is invalid"):
         api_tool.query({}, "")
 
 
-@patch("core.external_data_tool.api.api.db")
 @patch("core.external_data_tool.api.api.encrypter")
 @patch("core.external_data_tool.api.api.APIBasedExtensionRequestor")
-def test_query_requestor_init_error(mock_requestor_class, mock_encrypter, mock_db, api_tool):
-    mock_extension = MagicMock()
-    mock_extension.id = "ext_id"
-    mock_extension.tenant_id = "tenant_id"
-    mock_extension.api_endpoint = "http://api"
-    mock_extension.api_key = "encrypted_key"
-    mock_db.session.scalar.return_value = mock_extension
+def test_query_requestor_init_error(mock_requestor_class, mock_encrypter, api_tool, api_extension: APIBasedExtension):
     mock_encrypter.decrypt_token.return_value = "decrypted_key"
 
     mock_requestor_class.side_effect = Exception("init error")
@@ -107,16 +116,9 @@ def test_query_requestor_init_error(mock_requestor_class, mock_encrypter, mock_d
         api_tool.query({}, "")
 
 
-@patch("core.external_data_tool.api.api.db")
 @patch("core.external_data_tool.api.api.encrypter")
 @patch("core.external_data_tool.api.api.APIBasedExtensionRequestor")
-def test_query_no_result_in_response(mock_requestor_class, mock_encrypter, mock_db, api_tool):
-    mock_extension = MagicMock()
-    mock_extension.id = "ext_id"
-    mock_extension.tenant_id = "tenant_id"
-    mock_extension.api_endpoint = "http://api"
-    mock_extension.api_key = "encrypted_key"
-    mock_db.session.scalar.return_value = mock_extension
+def test_query_no_result_in_response(mock_requestor_class, mock_encrypter, api_tool, api_extension: APIBasedExtension):
     mock_encrypter.decrypt_token.return_value = "decrypted_key"
 
     mock_requestor = mock_requestor_class.return_value
@@ -126,16 +128,9 @@ def test_query_no_result_in_response(mock_requestor_class, mock_encrypter, mock_
         api_tool.query({}, "")
 
 
-@patch("core.external_data_tool.api.api.db")
 @patch("core.external_data_tool.api.api.encrypter")
 @patch("core.external_data_tool.api.api.APIBasedExtensionRequestor")
-def test_query_result_not_string(mock_requestor_class, mock_encrypter, mock_db, api_tool):
-    mock_extension = MagicMock()
-    mock_extension.id = "ext_id"
-    mock_extension.tenant_id = "tenant_id"
-    mock_extension.api_endpoint = "http://api"
-    mock_extension.api_key = "encrypted_key"
-    mock_db.session.scalar.return_value = mock_extension
+def test_query_result_not_string(mock_requestor_class, mock_encrypter, api_tool, api_extension: APIBasedExtension):
     mock_encrypter.decrypt_token.return_value = "decrypted_key"
 
     mock_requestor = mock_requestor_class.return_value