mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-09-24 15:45:26 +00:00
abe2706298
* Move filter utilities * Add MetadataRouter * Add tests for MetadataRouter * Add more tests * Rename FileExtensionClassifer to FileExtensionRouter * Add support for dates in filters * Add tests * Add release note * Add release note * Apply suggestions from code review Co-authored-by: ZanSara <sara.zanzottera@deepset.ai> --------- Co-authored-by: ZanSara <sara.zanzottera@deepset.ai>
100 lines
3.7 KiB
Python
100 lines
3.7 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
from haystack.preview.components.routers.file_router import FileExtensionRouter
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform in ["win32", "cygwin"],
|
|
reason="Can't run on Windows Github CI, need access to registry to get mime types",
|
|
)
|
|
class TestFileExtensionRouter:
|
|
@pytest.mark.unit
|
|
def test_to_dict(self):
|
|
component = FileExtensionRouter(mime_types=["text/plain", "audio/x-wav", "image/jpeg"])
|
|
data = component.to_dict()
|
|
assert data == {
|
|
"type": "FileExtensionRouter",
|
|
"init_parameters": {"mime_types": ["text/plain", "audio/x-wav", "image/jpeg"]},
|
|
}
|
|
|
|
@pytest.mark.unit
|
|
def test_from_dict(self):
|
|
data = {
|
|
"type": "FileExtensionRouter",
|
|
"init_parameters": {"mime_types": ["text/plain", "audio/x-wav", "image/jpeg"]},
|
|
}
|
|
component = FileExtensionRouter.from_dict(data)
|
|
assert component.mime_types == ["text/plain", "audio/x-wav", "image/jpeg"]
|
|
|
|
@pytest.mark.unit
|
|
def test_run(self, preview_samples_path):
|
|
"""
|
|
Test if the component runs correctly in the simplest happy path.
|
|
"""
|
|
file_paths = [
|
|
preview_samples_path / "txt" / "doc_1.txt",
|
|
preview_samples_path / "txt" / "doc_2.txt",
|
|
preview_samples_path / "audio" / "the context for this answer is here.wav",
|
|
preview_samples_path / "images" / "apple.jpg",
|
|
]
|
|
|
|
router = FileExtensionRouter(mime_types=["text/plain", "audio/x-wav", "image/jpeg"])
|
|
output = router.run(paths=file_paths)
|
|
assert output
|
|
assert len(output["text/plain"]) == 2
|
|
assert len(output["audio/x-wav"]) == 1
|
|
assert len(output["image/jpeg"]) == 1
|
|
assert not output["unclassified"]
|
|
|
|
@pytest.mark.unit
|
|
def test_no_files(self):
|
|
"""
|
|
Test that the component runs correctly when no files are provided.
|
|
"""
|
|
router = FileExtensionRouter(mime_types=["text/plain", "audio/x-wav", "image/jpeg"])
|
|
output = router.run(paths=[])
|
|
assert not output
|
|
|
|
@pytest.mark.unit
|
|
def test_unlisted_extensions(self, preview_samples_path):
|
|
"""
|
|
Test that the component correctly handles files with non specified mime types.
|
|
"""
|
|
file_paths = [
|
|
preview_samples_path / "txt" / "doc_1.txt",
|
|
preview_samples_path / "audio" / "ignored.mp3",
|
|
preview_samples_path / "audio" / "this is the content of the document.wav",
|
|
]
|
|
router = FileExtensionRouter(mime_types=["text/plain"])
|
|
output = router.run(paths=file_paths)
|
|
assert len(output["text/plain"]) == 1
|
|
assert "mp3" not in output
|
|
assert len(output["unclassified"]) == 2
|
|
assert str(output["unclassified"][0]).endswith("ignored.mp3")
|
|
assert str(output["unclassified"][1]).endswith("this is the content of the document.wav")
|
|
|
|
@pytest.mark.unit
|
|
def test_no_extension(self, preview_samples_path):
|
|
"""
|
|
Test that the component ignores files with no extension.
|
|
"""
|
|
file_paths = [
|
|
preview_samples_path / "txt" / "doc_1.txt",
|
|
preview_samples_path / "txt" / "doc_2",
|
|
preview_samples_path / "txt" / "doc_2.txt",
|
|
]
|
|
router = FileExtensionRouter(mime_types=["text/plain"])
|
|
output = router.run(paths=file_paths)
|
|
assert len(output["text/plain"]) == 2
|
|
assert len(output["unclassified"]) == 1
|
|
|
|
@pytest.mark.unit
|
|
def test_unknown_mime_type(self):
|
|
"""
|
|
Test that the component handles files with unknown mime types.
|
|
"""
|
|
with pytest.raises(ValueError, match="Unknown mime type:"):
|
|
FileExtensionRouter(mime_types=["type_invalid"])
|