From c3044b1bf03f0a5efaf755f12aab08cdb8cb1c8c Mon Sep 17 00:00:00 2001
From: Mike Wang <62768671+skcoirz@users.noreply.github.com>
Date: Fri, 5 May 2023 14:49:02 -0700
Subject: [PATCH] [test] Add integration_test for PandasAgent (#4056)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- confirm creation
- confirm functionality with a simple dimension check.

The test now is calling OpenAI API directly, but learning from
@vowelparrot that we’re caching the requests, so that it’s not that
expensive. I also found we’re calling OpenAI api in other integration
tests. Please lmk if there is any concern of real external API calls. I
can alternatively make a fake LLM for this test. Thanks
---
 .../agent/test_pandas_agent.py                | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)
 create mode 100644 tests/integration_tests/agent/test_pandas_agent.py

diff --git a/tests/integration_tests/agent/test_pandas_agent.py b/tests/integration_tests/agent/test_pandas_agent.py
new file mode 100644
index 00000000..9fd545e2
--- /dev/null
+++ b/tests/integration_tests/agent/test_pandas_agent.py
@@ -0,0 +1,30 @@
+import re
+
+import numpy as np
+import pytest
+from pandas import DataFrame
+
+from langchain.agents import create_pandas_dataframe_agent
+from langchain.agents.agent import AgentExecutor
+from langchain.llms import OpenAI
+
+
+@pytest.fixture(scope="module")
+def df() -> DataFrame:
+    random_data = np.random.rand(4, 4)
+    df = DataFrame(random_data, columns=["name", "age", "food", "sport"])
+    return df
+
+
+def test_pandas_agent_creation(df: DataFrame) -> None:
+    agent = create_pandas_dataframe_agent(OpenAI(temperature=0), df)
+    assert isinstance(agent, AgentExecutor)
+
+
+def test_data_reading(df: DataFrame) -> None:
+    agent = create_pandas_dataframe_agent(OpenAI(temperature=0), df)
+    assert isinstance(agent, AgentExecutor)
+    response = agent.run("how many rows in df? Give me a number.")
+    result = re.search(rf".*({df.shape[0]}).*", response)
+    assert result is not None
+    assert result.group(1) is not None