mirror of
https://github.com/hwchase17/langchain
synced 2024-11-10 01:10:59 +00:00
44a3484503
Thank you for contributing to LangChain! - **Description:** added unit tests for NotebookLoader. Linked PR: https://github.com/langchain-ai/langchain/pull/17614 - **Issue:** [#17614](https://github.com/langchain-ai/langchain/pull/17614) - **Twitter handle:** @paulodoestech - [x] Pass lint and test: Run `make format`, `make lint` and `make test` from the root of the package(s) you've modified to check that you're passing lint and testing. See contribution guidelines for more information on how to write/run tests, lint, etc: https://python.langchain.com/docs/contributing/ - [x] Add tests and docs: If you're adding a new integration, please include 1. a test for the integration, preferably unit tests that do not rely on network access, 2. an example notebook showing its use. It lives in `docs/docs/integrations` directory. If no one reviews your PR within a few days, please @-mention one of baskaryan, efriis, eyurtsev, hwchase17. --------- Co-authored-by: lachiewalker <lachiewalker1@hotmail.com> Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com> Co-authored-by: Bagatur <baskaryan@gmail.com>
86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
import json
|
|
|
|
from pytest_mock import MockerFixture
|
|
|
|
from langchain_community.document_loaders.notebook import NotebookLoader
|
|
|
|
|
|
def test_initialization() -> None:
|
|
loader = NotebookLoader(path="./testfile.ipynb")
|
|
assert loader.file_path == "./testfile.ipynb"
|
|
|
|
|
|
def test_load_no_outputs(mocker: MockerFixture) -> None:
|
|
mock_notebook_content = {
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": ["# Test notebook\n", "This is a test notebook."],
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": ["Hello World!\n"],
|
|
}
|
|
],
|
|
}
|
|
]
|
|
}
|
|
mocked_cell_type = mock_notebook_content["cells"][0]["cell_type"]
|
|
mocked_source = mock_notebook_content["cells"][0]["source"]
|
|
|
|
# Convert the mock notebook content to a JSON string
|
|
mock_notebook_content_str = json.dumps(mock_notebook_content)
|
|
|
|
# Mock the open function & json.load functions
|
|
mocker.patch("builtins.open", mocker.mock_open(read_data=mock_notebook_content_str))
|
|
mocker.patch("json.load", return_value=mock_notebook_content)
|
|
|
|
loader = NotebookLoader(path="./testfile.ipynb")
|
|
docs = loader.load()
|
|
|
|
assert len(docs) == 1
|
|
assert docs[0].page_content == f"'{mocked_cell_type}' cell: '{mocked_source}'\n\n"
|
|
assert docs[0].metadata == {"source": "testfile.ipynb"}
|
|
|
|
|
|
def test_load_with_outputs(mocker: MockerFixture) -> None:
|
|
mock_notebook_content: dict = {
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"metadata": {},
|
|
"source": ["# Test notebook\n", "This is a test notebook."],
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": ["Hello World!\n"],
|
|
}
|
|
],
|
|
}
|
|
]
|
|
}
|
|
mocked_cell_type = mock_notebook_content["cells"][0]["cell_type"]
|
|
mocked_source = mock_notebook_content["cells"][0]["source"]
|
|
mocked_output = mock_notebook_content["cells"][0]["outputs"][0]["text"]
|
|
|
|
# Convert the mock notebook content to a JSON string
|
|
mock_notebook_content_str = json.dumps(mock_notebook_content)
|
|
|
|
# Mock the open function & json.load functions
|
|
mocker.patch("builtins.open", mocker.mock_open(read_data=mock_notebook_content_str))
|
|
mocker.patch("json.load", return_value=mock_notebook_content)
|
|
|
|
loader = NotebookLoader(path="./testfile.ipynb", include_outputs=True)
|
|
docs = loader.load()
|
|
|
|
assert len(docs) == 1
|
|
expected_content = (
|
|
f"'{mocked_cell_type}' cell: '{mocked_source}'\n"
|
|
f" with output: '{mocked_output}'\n\n"
|
|
)
|
|
assert docs[0].page_content == expected_content
|
|
assert docs[0].metadata == {"source": "testfile.ipynb"}
|