class_docstr: Unit test for Smart Data Lake Prompt Generation
class_docstr_tok: ["Unit","test","for","Smart","Data","Lake","Prompt","Generation"]
code: def test_get_chat_prompt(self, context):
      # Test case 1: direct_sql is True
      prompt_generation = PromptGeneration()
      context.config.direct_sql = True
  
      gen_key, gen_prompt = prompt_generation._get_chat_prompt(context)
      expected_key = "direct_sql_prompt"
      assert gen_key == expected_key
      assert isinstance(gen_prompt, DirectSQLPrompt)
  
      # Test case 2: direct_sql is False
      context.config.direct_sql = False
  
      gen_key, gen_prompt = prompt_generation._get_chat_prompt(context)
      expected_key = "generate_python_code"
      assert gen_key == expected_key
      assert isinstance(gen_prompt, GeneratePythonCodePrompt)
  
file_path: pandas-ai/tests/pipelines/smart_datalake/test_prompt_generation.py
imports: import pandas
  from typing import Optional
  import pytest
  from pandasai.llm.fake import FakeLLM
  from pandasai.prompts.direct_sql_prompt import DirectSQLPrompt
  from pandasai.prompts.generate_python_code import GeneratePythonCodePrompt
  from pandasai.pipelines.smart_datalake_chat.prompt_generation import PromptGeneration
  from pandasai.pipelines.pipeline_context import PipelineContext
label_desc: Use image node_id 6 for calling the TestPromptGeneration obj's underlying member method code with example usage: obj.test_get_chat_prompt(context) without return types