diff --git a/python-threatexchange/threatexchange/content_type/tests/test_file_content_type.py b/python-threatexchange/threatexchange/content_type/tests/test_file_content_type.py new file mode 100644 index 000000000..077052353 --- /dev/null +++ b/python-threatexchange/threatexchange/content_type/tests/test_file_content_type.py @@ -0,0 +1,27 @@ +import pytest +from threatexchange.content_type.photo import PhotoContent +from threatexchange.content_type.video import VideoContent +from threatexchange.content_type.file_content import FileContent + +@pytest.mark.parametrize("file_name,expected_content_type", [ + ("file.jpg", PhotoContent), + ("file.JPG", PhotoContent), + ("file.mp4", VideoContent), + ("file.MP4", VideoContent), + ("archive.photo.png", PhotoContent), + ("movie.backup.mp4", VideoContent), +]) +def test_file_content_detection(file_name, expected_content_type): + """ + Tests that FileContent correctly identifies the content type + as either PhotoContent or VideoContent based on file extension. + """ + content_type = FileContent.get_content_type_from_filename(file_name) + assert content_type == expected_content_type, f"Failed for {file_name}" + +def test_unknown_file_type(): + """ + Tests that an unknown file type returns None. + """ + file_content = FileContent.get_content_type_from_filename("file.txt") + assert file_content is None diff --git a/python-threatexchange/threatexchange/tests/test_file_content_type.py b/python-threatexchange/threatexchange/tests/test_file_content_type.py deleted file mode 100644 index ba27094cc..000000000 --- a/python-threatexchange/threatexchange/tests/test_file_content_type.py +++ /dev/null @@ -1,33 +0,0 @@ -import unittest -from threatexchange.content_type.photo import PhotoContent -from threatexchange.content_type.video import VideoContent -from threatexchange.content_type.file_content import FileContent - -class TestFileContentType(unittest.TestCase): - def test_photo_detection_jpg(self): - file_content = FileContent.get_content_type_from_filename("file.jpg") - self.assertEqual(file_content, PhotoContent) - - def test_photo_detection_uppercase_extension(self): - file_content = FileContent.get_content_type_from_filename("file.JPG") - self.assertEqual(file_content, PhotoContent) - - def test_video_detection_mp4(self): - file_content = FileContent.get_content_type_from_filename("file.mp4") - self.assertEqual(file_content, VideoContent) - - def test_video_detection_uppercase_extension(self): - file_content = FileContent.get_content_type_from_filename("file.MP4") - self.assertEqual(file_content, VideoContent) - - def test_unknown_file_type(self): - file_content = FileContent.get_content_type_from_filename("file.txt") - self.assertIsNone(file_content) - - def test_photo_with_multiple_dots(self): - file_content = FileContent.get_content_type_from_filename("archive.photo.png") - self.assertEqual(file_content, PhotoContent) - - def test_video_with_multiple_dots(self): - file_content = FileContent.get_content_type_from_filename("movie.backup.mp4") - self.assertEqual(file_content, VideoContent) \ No newline at end of file