class_docstr: Unit tests for the CodeManager class
class_docstr_tok: ["Unit","tests","for","the","CodeManager","class"]
docstr: Test that the direct SQL function definition is removed when 'direct_sql' is True
docstr_tok: ["Test","that","the","direct","SQL","function","definition","is","removed","when","'direct_sql","'","is","True"]
code: def test_clean_code_direct_sql_code(
      self,
      pgsql_connector: PostgreSQLConnector,
      exec_context: MagicMock,
  ):
      """Test that the direct SQL function definition is removed when 'direct_sql' is True"""
      df = SmartDataframe(
          pgsql_connector,
          config={"llm": FakeLLM(output=""), "direct_sql": True},
      )
      code_manager = df.lake._code_manager
      safe_code = """
  import numpy as np
  def execute_sql_query(sql_query: str) -> pd.DataFrame:
      # code to connect to the database and execute the query
      # ...
      # return the result as a dataframe
      return pd.DataFrame()
  np.array()
  """
      assert (
          code_manager._clean_code(safe_code, exec_context)
          == "np.array()"
      )
  
imports: import ast
  from typing import Optional
  from unittest.mock import MagicMock, Mock, patch
  import uuid
  import pandas
  import pytest
  from pandasai.connectors.base import SQLConnectorConfig
  from pandasai.connectors.sql import PostgreSQLConnector, SQLConnector
  from pandasai.exceptions import BadImportError, MaliciousQueryError, NoCodeFoundError
  from pandasai.helpers.skills_manager import SkillsManager
  from pandasai.llm.fake import FakeLLM
  from pandasai.exceptions import InvalidConfigError
  from pandasai.smart_dataframe import SmartDataframe
  from pandasai.helpers.code_manager import CodeExecutionContext, CodeManager
label_desc: Use image node_id 25 for calling the TestCodeManager obj's underlying member method code with example usage: obj.test_clean_code_direct_sql_code(pgsql_connector, exec_context) without return types