mirror of https://github.com/hwchase17/langchain
community[patch]: add NotebookLoader unit test (#17721)
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>pull/17614/head^2
parent
4c3a67122f
commit
44a3484503
@ -0,0 +1,85 @@
|
||||
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"}
|
Loading…
Reference in New Issue