haystack/test/preview/components/routers/test_file_router.py
bogdankostic abe2706298
feat: Add MetadataRouter (2.0) (#5824)
* 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>
2023-09-20 14:49:17 +02:00

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"])