mirror of
https://github.com/zylon-ai/private-gpt.git
synced 2025-12-22 13:55:41 +01:00
68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
import tempfile
|
|
from pathlib import Path
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from private_gpt.server.ingest.ingest_router import IngestResponse
|
|
from tests.fixtures.ingest_helper import IngestHelper
|
|
|
|
|
|
def test_ingest_accepts_txt_files(ingest_helper: IngestHelper) -> None:
|
|
path = Path(__file__).parents[0] / "test.txt"
|
|
ingest_result = ingest_helper.ingest_file(path)
|
|
assert len(ingest_result.data) == 1
|
|
|
|
|
|
def test_ingest_accepts_pdf_files(ingest_helper: IngestHelper) -> None:
|
|
path = Path(__file__).parents[0] / "test.pdf"
|
|
ingest_result = ingest_helper.ingest_file(path)
|
|
assert len(ingest_result.data) == 1
|
|
|
|
|
|
def test_ingest_list_returns_something_after_ingestion(
|
|
test_client: TestClient, ingest_helper: IngestHelper
|
|
) -> None:
|
|
response_before = test_client.get("/v1/ingest/list")
|
|
count_ingest_before = len(response_before.json()["data"])
|
|
with tempfile.NamedTemporaryFile("w", suffix=".txt") as test_file:
|
|
test_file.write("Foo bar; hello there!")
|
|
test_file.flush()
|
|
test_file.seek(0)
|
|
ingest_result = ingest_helper.ingest_file(Path(test_file.name))
|
|
assert len(ingest_result.data) == 1, "The temp doc should have been ingested"
|
|
response_after = test_client.get("/v1/ingest/list")
|
|
count_ingest_after = len(response_after.json()["data"])
|
|
assert (
|
|
count_ingest_after == count_ingest_before + 1
|
|
), "The temp doc should be returned"
|
|
|
|
|
|
def test_ingest_plain_text(test_client: TestClient) -> None:
|
|
response = test_client.post(
|
|
"/v1/ingest/text", json={"file_name": "file_name", "text": "text"}
|
|
)
|
|
assert response.status_code == 200
|
|
ingest_result = IngestResponse.model_validate(response.json())
|
|
assert len(ingest_result.data) == 1
|
|
|
|
|
|
def test_ingest_text_with_metadata(test_client: TestClient):
|
|
response = test_client.post(
|
|
"/v1/ingest/text",
|
|
json={"file_name": "file_name", "text": "text", "metadata": {"foo": "bar"}},
|
|
)
|
|
assert response.status_code == 200
|
|
ingest_result = IngestResponse.model_validate(response.json())
|
|
assert len(ingest_result.data) == 1
|
|
|
|
assert ingest_result.data[0].doc_metadata == {
|
|
"file_name": "file_name",
|
|
"foo": "bar",
|
|
}
|
|
|
|
|
|
def test_ingest_accepts_txt_files_with_metadata(ingest_helper: IngestHelper) -> None:
|
|
path = Path(__file__).parents[0] / "test.txt"
|
|
ingest_result = ingest_helper.ingest_file_with_metadata(path, {"foo": "bar"})
|
|
assert len(ingest_result.data) == 1
|
|
assert ingest_result.data[0].doc_metadata == {"file_name": "test.txt", "foo": "bar"}
|