test_file_service.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913
  1. import hashlib
  2. from io import BytesIO
  3. from unittest.mock import patch
  4. import pytest
  5. from faker import Faker
  6. from werkzeug.exceptions import NotFound
  7. from configs import dify_config
  8. from models.account import Account, Tenant
  9. from models.enums import CreatorUserRole
  10. from models.model import EndUser, UploadFile
  11. from services.errors.file import FileTooLargeError, UnsupportedFileTypeError
  12. from services.file_service import FileService
  13. class TestFileService:
  14. """Integration tests for FileService using testcontainers."""
  15. @pytest.fixture
  16. def mock_external_service_dependencies(self):
  17. """Mock setup for external service dependencies."""
  18. with (
  19. patch("services.file_service.storage") as mock_storage,
  20. patch("services.file_service.file_helpers") as mock_file_helpers,
  21. patch("services.file_service.ExtractProcessor") as mock_extract_processor,
  22. ):
  23. # Setup default mock returns
  24. mock_storage.save.return_value = None
  25. mock_storage.load.return_value = BytesIO(b"mock file content")
  26. mock_file_helpers.get_signed_file_url.return_value = "https://example.com/signed-url"
  27. mock_file_helpers.verify_image_signature.return_value = True
  28. mock_file_helpers.verify_file_signature.return_value = True
  29. mock_extract_processor.load_from_upload_file.return_value = "extracted text content"
  30. yield {
  31. "storage": mock_storage,
  32. "file_helpers": mock_file_helpers,
  33. "extract_processor": mock_extract_processor,
  34. }
  35. def _create_test_account(self, db_session_with_containers, mock_external_service_dependencies):
  36. """
  37. Helper method to create a test account for testing.
  38. Args:
  39. db_session_with_containers: Database session from testcontainers infrastructure
  40. mock_external_service_dependencies: Mock dependencies
  41. Returns:
  42. Account: Created account instance
  43. """
  44. fake = Faker()
  45. # Create account
  46. account = Account(
  47. email=fake.email(),
  48. name=fake.name(),
  49. interface_language="en-US",
  50. status="active",
  51. )
  52. from extensions.ext_database import db
  53. db.session.add(account)
  54. db.session.commit()
  55. # Create tenant for the account
  56. tenant = Tenant(
  57. name=fake.company(),
  58. status="normal",
  59. )
  60. db.session.add(tenant)
  61. db.session.commit()
  62. # Create tenant-account join
  63. from models.account import TenantAccountJoin, TenantAccountRole
  64. join = TenantAccountJoin(
  65. tenant_id=tenant.id,
  66. account_id=account.id,
  67. role=TenantAccountRole.OWNER.value,
  68. current=True,
  69. )
  70. db.session.add(join)
  71. db.session.commit()
  72. # Set current tenant for account
  73. account.current_tenant = tenant
  74. return account
  75. def _create_test_end_user(self, db_session_with_containers, mock_external_service_dependencies):
  76. """
  77. Helper method to create a test end user for testing.
  78. Args:
  79. db_session_with_containers: Database session from testcontainers infrastructure
  80. mock_external_service_dependencies: Mock dependencies
  81. Returns:
  82. EndUser: Created end user instance
  83. """
  84. fake = Faker()
  85. end_user = EndUser(
  86. tenant_id=str(fake.uuid4()),
  87. type="web",
  88. name=fake.name(),
  89. is_anonymous=False,
  90. session_id=fake.uuid4(),
  91. )
  92. from extensions.ext_database import db
  93. db.session.add(end_user)
  94. db.session.commit()
  95. return end_user
  96. def _create_test_upload_file(self, db_session_with_containers, mock_external_service_dependencies, account):
  97. """
  98. Helper method to create a test upload file for testing.
  99. Args:
  100. db_session_with_containers: Database session from testcontainers infrastructure
  101. mock_external_service_dependencies: Mock dependencies
  102. account: Account instance
  103. Returns:
  104. UploadFile: Created upload file instance
  105. """
  106. fake = Faker()
  107. upload_file = UploadFile(
  108. tenant_id=account.current_tenant_id if hasattr(account, "current_tenant_id") else str(fake.uuid4()),
  109. storage_type="local",
  110. key=f"upload_files/test/{fake.uuid4()}.txt",
  111. name="test_file.txt",
  112. size=1024,
  113. extension="txt",
  114. mime_type="text/plain",
  115. created_by_role=CreatorUserRole.ACCOUNT,
  116. created_by=account.id,
  117. created_at=fake.date_time(),
  118. used=False,
  119. hash=hashlib.sha3_256(b"test content").hexdigest(),
  120. source_url="",
  121. )
  122. from extensions.ext_database import db
  123. db.session.add(upload_file)
  124. db.session.commit()
  125. return upload_file
  126. # Test upload_file method
  127. def test_upload_file_success(self, db_session_with_containers, mock_external_service_dependencies):
  128. """
  129. Test successful file upload with valid parameters.
  130. """
  131. fake = Faker()
  132. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  133. filename = "test_document.pdf"
  134. content = b"test file content"
  135. mimetype = "application/pdf"
  136. upload_file = FileService.upload_file(
  137. filename=filename,
  138. content=content,
  139. mimetype=mimetype,
  140. user=account,
  141. )
  142. assert upload_file is not None
  143. assert upload_file.name == filename
  144. assert upload_file.size == len(content)
  145. assert upload_file.extension == "pdf"
  146. assert upload_file.mime_type == mimetype
  147. assert upload_file.created_by == account.id
  148. assert upload_file.created_by_role == CreatorUserRole.ACCOUNT.value
  149. assert upload_file.used is False
  150. assert upload_file.hash == hashlib.sha3_256(content).hexdigest()
  151. # Verify storage was called
  152. mock_external_service_dependencies["storage"].save.assert_called_once()
  153. # Verify database state
  154. from extensions.ext_database import db
  155. db.session.refresh(upload_file)
  156. assert upload_file.id is not None
  157. def test_upload_file_with_end_user(self, db_session_with_containers, mock_external_service_dependencies):
  158. """
  159. Test file upload with end user instead of account.
  160. """
  161. fake = Faker()
  162. end_user = self._create_test_end_user(db_session_with_containers, mock_external_service_dependencies)
  163. filename = "test_image.jpg"
  164. content = b"test image content"
  165. mimetype = "image/jpeg"
  166. upload_file = FileService.upload_file(
  167. filename=filename,
  168. content=content,
  169. mimetype=mimetype,
  170. user=end_user,
  171. )
  172. assert upload_file is not None
  173. assert upload_file.created_by == end_user.id
  174. assert upload_file.created_by_role == CreatorUserRole.END_USER.value
  175. def test_upload_file_with_datasets_source(self, db_session_with_containers, mock_external_service_dependencies):
  176. """
  177. Test file upload with datasets source parameter.
  178. """
  179. fake = Faker()
  180. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  181. filename = "test_document.pdf"
  182. content = b"test file content"
  183. mimetype = "application/pdf"
  184. upload_file = FileService.upload_file(
  185. filename=filename,
  186. content=content,
  187. mimetype=mimetype,
  188. user=account,
  189. source="datasets",
  190. source_url="https://example.com/source",
  191. )
  192. assert upload_file is not None
  193. assert upload_file.source_url == "https://example.com/source"
  194. def test_upload_file_invalid_filename_characters(
  195. self, db_session_with_containers, mock_external_service_dependencies
  196. ):
  197. """
  198. Test file upload with invalid filename characters.
  199. """
  200. fake = Faker()
  201. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  202. filename = "test/file<name>.txt"
  203. content = b"test content"
  204. mimetype = "text/plain"
  205. with pytest.raises(ValueError, match="Filename contains invalid characters"):
  206. FileService.upload_file(
  207. filename=filename,
  208. content=content,
  209. mimetype=mimetype,
  210. user=account,
  211. )
  212. def test_upload_file_filename_too_long(self, db_session_with_containers, mock_external_service_dependencies):
  213. """
  214. Test file upload with filename that exceeds length limit.
  215. """
  216. fake = Faker()
  217. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  218. # Create a filename longer than 200 characters
  219. long_name = "a" * 250
  220. filename = f"{long_name}.txt"
  221. content = b"test content"
  222. mimetype = "text/plain"
  223. upload_file = FileService.upload_file(
  224. filename=filename,
  225. content=content,
  226. mimetype=mimetype,
  227. user=account,
  228. )
  229. # Verify filename was truncated (the logic truncates the base name to 200 chars + extension)
  230. # So the total length should be <= 200 + len(extension) + 1 (for the dot)
  231. assert len(upload_file.name) <= 200 + len(upload_file.extension) + 1
  232. assert upload_file.name.endswith(".txt")
  233. # Verify the base name was truncated
  234. base_name = upload_file.name[:-4] # Remove .txt
  235. assert len(base_name) <= 200
  236. def test_upload_file_datasets_unsupported_type(
  237. self, db_session_with_containers, mock_external_service_dependencies
  238. ):
  239. """
  240. Test file upload for datasets with unsupported file type.
  241. """
  242. fake = Faker()
  243. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  244. filename = "test_image.jpg"
  245. content = b"test content"
  246. mimetype = "image/jpeg"
  247. with pytest.raises(UnsupportedFileTypeError):
  248. FileService.upload_file(
  249. filename=filename,
  250. content=content,
  251. mimetype=mimetype,
  252. user=account,
  253. source="datasets",
  254. )
  255. def test_upload_file_too_large(self, db_session_with_containers, mock_external_service_dependencies):
  256. """
  257. Test file upload with file size exceeding limit.
  258. """
  259. fake = Faker()
  260. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  261. filename = "large_image.jpg"
  262. # Create content larger than the limit
  263. content = b"x" * (dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT * 1024 * 1024 + 1)
  264. mimetype = "image/jpeg"
  265. with pytest.raises(FileTooLargeError):
  266. FileService.upload_file(
  267. filename=filename,
  268. content=content,
  269. mimetype=mimetype,
  270. user=account,
  271. )
  272. # Test is_file_size_within_limit method
  273. def test_is_file_size_within_limit_image_success(
  274. self, db_session_with_containers, mock_external_service_dependencies
  275. ):
  276. """
  277. Test file size check for image files within limit.
  278. """
  279. extension = "jpg"
  280. file_size = dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT * 1024 * 1024 # Exactly at limit
  281. result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
  282. assert result is True
  283. def test_is_file_size_within_limit_video_success(
  284. self, db_session_with_containers, mock_external_service_dependencies
  285. ):
  286. """
  287. Test file size check for video files within limit.
  288. """
  289. extension = "mp4"
  290. file_size = dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT * 1024 * 1024 # Exactly at limit
  291. result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
  292. assert result is True
  293. def test_is_file_size_within_limit_audio_success(
  294. self, db_session_with_containers, mock_external_service_dependencies
  295. ):
  296. """
  297. Test file size check for audio files within limit.
  298. """
  299. extension = "mp3"
  300. file_size = dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT * 1024 * 1024 # Exactly at limit
  301. result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
  302. assert result is True
  303. def test_is_file_size_within_limit_document_success(
  304. self, db_session_with_containers, mock_external_service_dependencies
  305. ):
  306. """
  307. Test file size check for document files within limit.
  308. """
  309. extension = "pdf"
  310. file_size = dify_config.UPLOAD_FILE_SIZE_LIMIT * 1024 * 1024 # Exactly at limit
  311. result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
  312. assert result is True
  313. def test_is_file_size_within_limit_image_exceeded(
  314. self, db_session_with_containers, mock_external_service_dependencies
  315. ):
  316. """
  317. Test file size check for image files exceeding limit.
  318. """
  319. extension = "jpg"
  320. file_size = dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT * 1024 * 1024 + 1 # Exceeds limit
  321. result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
  322. assert result is False
  323. def test_is_file_size_within_limit_unknown_extension(
  324. self, db_session_with_containers, mock_external_service_dependencies
  325. ):
  326. """
  327. Test file size check for unknown file extension.
  328. """
  329. extension = "xyz"
  330. file_size = dify_config.UPLOAD_FILE_SIZE_LIMIT * 1024 * 1024 # Uses default limit
  331. result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
  332. assert result is True
  333. # Test upload_text method
  334. def test_upload_text_success(self, db_session_with_containers, mock_external_service_dependencies):
  335. """
  336. Test successful text upload.
  337. """
  338. fake = Faker()
  339. text = "This is a test text content"
  340. text_name = "test_text.txt"
  341. # Mock current_user
  342. with patch("services.file_service.current_user") as mock_current_user:
  343. mock_current_user.current_tenant_id = str(fake.uuid4())
  344. mock_current_user.id = str(fake.uuid4())
  345. upload_file = FileService.upload_text(text=text, text_name=text_name)
  346. assert upload_file is not None
  347. assert upload_file.name == text_name
  348. assert upload_file.size == len(text)
  349. assert upload_file.extension == "txt"
  350. assert upload_file.mime_type == "text/plain"
  351. assert upload_file.used is True
  352. assert upload_file.used_by == mock_current_user.id
  353. # Verify storage was called
  354. mock_external_service_dependencies["storage"].save.assert_called_once()
  355. def test_upload_text_name_too_long(self, db_session_with_containers, mock_external_service_dependencies):
  356. """
  357. Test text upload with name that exceeds length limit.
  358. """
  359. fake = Faker()
  360. text = "test content"
  361. long_name = "a" * 250 # Longer than 200 characters
  362. # Mock current_user
  363. with patch("services.file_service.current_user") as mock_current_user:
  364. mock_current_user.current_tenant_id = str(fake.uuid4())
  365. mock_current_user.id = str(fake.uuid4())
  366. upload_file = FileService.upload_text(text=text, text_name=long_name)
  367. # Verify name was truncated
  368. assert len(upload_file.name) <= 200
  369. assert upload_file.name == "a" * 200
  370. # Test get_file_preview method
  371. def test_get_file_preview_success(self, db_session_with_containers, mock_external_service_dependencies):
  372. """
  373. Test successful file preview generation.
  374. """
  375. fake = Faker()
  376. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  377. upload_file = self._create_test_upload_file(
  378. db_session_with_containers, mock_external_service_dependencies, account
  379. )
  380. # Update file to have document extension
  381. upload_file.extension = "pdf"
  382. from extensions.ext_database import db
  383. db.session.commit()
  384. result = FileService.get_file_preview(file_id=upload_file.id)
  385. assert result == "extracted text content"
  386. mock_external_service_dependencies["extract_processor"].load_from_upload_file.assert_called_once()
  387. def test_get_file_preview_file_not_found(self, db_session_with_containers, mock_external_service_dependencies):
  388. """
  389. Test file preview with non-existent file.
  390. """
  391. fake = Faker()
  392. non_existent_id = str(fake.uuid4())
  393. with pytest.raises(NotFound, match="File not found"):
  394. FileService.get_file_preview(file_id=non_existent_id)
  395. def test_get_file_preview_unsupported_file_type(
  396. self, db_session_with_containers, mock_external_service_dependencies
  397. ):
  398. """
  399. Test file preview with unsupported file type.
  400. """
  401. fake = Faker()
  402. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  403. upload_file = self._create_test_upload_file(
  404. db_session_with_containers, mock_external_service_dependencies, account
  405. )
  406. # Update file to have non-document extension
  407. upload_file.extension = "jpg"
  408. from extensions.ext_database import db
  409. db.session.commit()
  410. with pytest.raises(UnsupportedFileTypeError):
  411. FileService.get_file_preview(file_id=upload_file.id)
  412. def test_get_file_preview_text_truncation(self, db_session_with_containers, mock_external_service_dependencies):
  413. """
  414. Test file preview with text that exceeds preview limit.
  415. """
  416. fake = Faker()
  417. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  418. upload_file = self._create_test_upload_file(
  419. db_session_with_containers, mock_external_service_dependencies, account
  420. )
  421. # Update file to have document extension
  422. upload_file.extension = "pdf"
  423. from extensions.ext_database import db
  424. db.session.commit()
  425. # Mock long text content
  426. long_text = "x" * 5000 # Longer than PREVIEW_WORDS_LIMIT
  427. mock_external_service_dependencies["extract_processor"].load_from_upload_file.return_value = long_text
  428. result = FileService.get_file_preview(file_id=upload_file.id)
  429. assert len(result) == 3000 # PREVIEW_WORDS_LIMIT
  430. assert result == "x" * 3000
  431. # Test get_image_preview method
  432. def test_get_image_preview_success(self, db_session_with_containers, mock_external_service_dependencies):
  433. """
  434. Test successful image preview generation.
  435. """
  436. fake = Faker()
  437. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  438. upload_file = self._create_test_upload_file(
  439. db_session_with_containers, mock_external_service_dependencies, account
  440. )
  441. # Update file to have image extension
  442. upload_file.extension = "jpg"
  443. from extensions.ext_database import db
  444. db.session.commit()
  445. timestamp = "1234567890"
  446. nonce = "test_nonce"
  447. sign = "test_signature"
  448. generator, mime_type = FileService.get_image_preview(
  449. file_id=upload_file.id,
  450. timestamp=timestamp,
  451. nonce=nonce,
  452. sign=sign,
  453. )
  454. assert generator is not None
  455. assert mime_type == upload_file.mime_type
  456. mock_external_service_dependencies["file_helpers"].verify_image_signature.assert_called_once()
  457. def test_get_image_preview_invalid_signature(self, db_session_with_containers, mock_external_service_dependencies):
  458. """
  459. Test image preview with invalid signature.
  460. """
  461. fake = Faker()
  462. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  463. upload_file = self._create_test_upload_file(
  464. db_session_with_containers, mock_external_service_dependencies, account
  465. )
  466. # Mock invalid signature
  467. mock_external_service_dependencies["file_helpers"].verify_image_signature.return_value = False
  468. timestamp = "1234567890"
  469. nonce = "test_nonce"
  470. sign = "invalid_signature"
  471. with pytest.raises(NotFound, match="File not found or signature is invalid"):
  472. FileService.get_image_preview(
  473. file_id=upload_file.id,
  474. timestamp=timestamp,
  475. nonce=nonce,
  476. sign=sign,
  477. )
  478. def test_get_image_preview_file_not_found(self, db_session_with_containers, mock_external_service_dependencies):
  479. """
  480. Test image preview with non-existent file.
  481. """
  482. fake = Faker()
  483. non_existent_id = str(fake.uuid4())
  484. timestamp = "1234567890"
  485. nonce = "test_nonce"
  486. sign = "test_signature"
  487. with pytest.raises(NotFound, match="File not found or signature is invalid"):
  488. FileService.get_image_preview(
  489. file_id=non_existent_id,
  490. timestamp=timestamp,
  491. nonce=nonce,
  492. sign=sign,
  493. )
  494. def test_get_image_preview_unsupported_file_type(
  495. self, db_session_with_containers, mock_external_service_dependencies
  496. ):
  497. """
  498. Test image preview with non-image file type.
  499. """
  500. fake = Faker()
  501. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  502. upload_file = self._create_test_upload_file(
  503. db_session_with_containers, mock_external_service_dependencies, account
  504. )
  505. # Update file to have non-image extension
  506. upload_file.extension = "pdf"
  507. from extensions.ext_database import db
  508. db.session.commit()
  509. timestamp = "1234567890"
  510. nonce = "test_nonce"
  511. sign = "test_signature"
  512. with pytest.raises(UnsupportedFileTypeError):
  513. FileService.get_image_preview(
  514. file_id=upload_file.id,
  515. timestamp=timestamp,
  516. nonce=nonce,
  517. sign=sign,
  518. )
  519. # Test get_file_generator_by_file_id method
  520. def test_get_file_generator_by_file_id_success(
  521. self, db_session_with_containers, mock_external_service_dependencies
  522. ):
  523. """
  524. Test successful file generator retrieval.
  525. """
  526. fake = Faker()
  527. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  528. upload_file = self._create_test_upload_file(
  529. db_session_with_containers, mock_external_service_dependencies, account
  530. )
  531. timestamp = "1234567890"
  532. nonce = "test_nonce"
  533. sign = "test_signature"
  534. generator, file_obj = FileService.get_file_generator_by_file_id(
  535. file_id=upload_file.id,
  536. timestamp=timestamp,
  537. nonce=nonce,
  538. sign=sign,
  539. )
  540. assert generator is not None
  541. assert file_obj == upload_file
  542. mock_external_service_dependencies["file_helpers"].verify_file_signature.assert_called_once()
  543. def test_get_file_generator_by_file_id_invalid_signature(
  544. self, db_session_with_containers, mock_external_service_dependencies
  545. ):
  546. """
  547. Test file generator retrieval with invalid signature.
  548. """
  549. fake = Faker()
  550. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  551. upload_file = self._create_test_upload_file(
  552. db_session_with_containers, mock_external_service_dependencies, account
  553. )
  554. # Mock invalid signature
  555. mock_external_service_dependencies["file_helpers"].verify_file_signature.return_value = False
  556. timestamp = "1234567890"
  557. nonce = "test_nonce"
  558. sign = "invalid_signature"
  559. with pytest.raises(NotFound, match="File not found or signature is invalid"):
  560. FileService.get_file_generator_by_file_id(
  561. file_id=upload_file.id,
  562. timestamp=timestamp,
  563. nonce=nonce,
  564. sign=sign,
  565. )
  566. def test_get_file_generator_by_file_id_file_not_found(
  567. self, db_session_with_containers, mock_external_service_dependencies
  568. ):
  569. """
  570. Test file generator retrieval with non-existent file.
  571. """
  572. fake = Faker()
  573. non_existent_id = str(fake.uuid4())
  574. timestamp = "1234567890"
  575. nonce = "test_nonce"
  576. sign = "test_signature"
  577. with pytest.raises(NotFound, match="File not found or signature is invalid"):
  578. FileService.get_file_generator_by_file_id(
  579. file_id=non_existent_id,
  580. timestamp=timestamp,
  581. nonce=nonce,
  582. sign=sign,
  583. )
  584. # Test get_public_image_preview method
  585. def test_get_public_image_preview_success(self, db_session_with_containers, mock_external_service_dependencies):
  586. """
  587. Test successful public image preview generation.
  588. """
  589. fake = Faker()
  590. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  591. upload_file = self._create_test_upload_file(
  592. db_session_with_containers, mock_external_service_dependencies, account
  593. )
  594. # Update file to have image extension
  595. upload_file.extension = "jpg"
  596. from extensions.ext_database import db
  597. db.session.commit()
  598. generator, mime_type = FileService.get_public_image_preview(file_id=upload_file.id)
  599. assert generator is not None
  600. assert mime_type == upload_file.mime_type
  601. mock_external_service_dependencies["storage"].load.assert_called_once()
  602. def test_get_public_image_preview_file_not_found(
  603. self, db_session_with_containers, mock_external_service_dependencies
  604. ):
  605. """
  606. Test public image preview with non-existent file.
  607. """
  608. fake = Faker()
  609. non_existent_id = str(fake.uuid4())
  610. with pytest.raises(NotFound, match="File not found or signature is invalid"):
  611. FileService.get_public_image_preview(file_id=non_existent_id)
  612. def test_get_public_image_preview_unsupported_file_type(
  613. self, db_session_with_containers, mock_external_service_dependencies
  614. ):
  615. """
  616. Test public image preview with non-image file type.
  617. """
  618. fake = Faker()
  619. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  620. upload_file = self._create_test_upload_file(
  621. db_session_with_containers, mock_external_service_dependencies, account
  622. )
  623. # Update file to have non-image extension
  624. upload_file.extension = "pdf"
  625. from extensions.ext_database import db
  626. db.session.commit()
  627. with pytest.raises(UnsupportedFileTypeError):
  628. FileService.get_public_image_preview(file_id=upload_file.id)
  629. # Test edge cases and boundary conditions
  630. def test_upload_file_empty_content(self, db_session_with_containers, mock_external_service_dependencies):
  631. """
  632. Test file upload with empty content.
  633. """
  634. fake = Faker()
  635. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  636. filename = "empty.txt"
  637. content = b""
  638. mimetype = "text/plain"
  639. upload_file = FileService.upload_file(
  640. filename=filename,
  641. content=content,
  642. mimetype=mimetype,
  643. user=account,
  644. )
  645. assert upload_file is not None
  646. assert upload_file.size == 0
  647. def test_upload_file_special_characters_in_name(
  648. self, db_session_with_containers, mock_external_service_dependencies
  649. ):
  650. """
  651. Test file upload with special characters in filename (but valid ones).
  652. """
  653. fake = Faker()
  654. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  655. filename = "test-file_with_underscores_and.dots.txt"
  656. content = b"test content"
  657. mimetype = "text/plain"
  658. upload_file = FileService.upload_file(
  659. filename=filename,
  660. content=content,
  661. mimetype=mimetype,
  662. user=account,
  663. )
  664. assert upload_file is not None
  665. assert upload_file.name == filename
  666. def test_upload_file_different_case_extensions(
  667. self, db_session_with_containers, mock_external_service_dependencies
  668. ):
  669. """
  670. Test file upload with different case extensions.
  671. """
  672. fake = Faker()
  673. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  674. filename = "test.PDF"
  675. content = b"test content"
  676. mimetype = "application/pdf"
  677. upload_file = FileService.upload_file(
  678. filename=filename,
  679. content=content,
  680. mimetype=mimetype,
  681. user=account,
  682. )
  683. assert upload_file is not None
  684. assert upload_file.extension == "pdf" # Should be converted to lowercase
  685. def test_upload_text_empty_text(self, db_session_with_containers, mock_external_service_dependencies):
  686. """
  687. Test text upload with empty text.
  688. """
  689. fake = Faker()
  690. text = ""
  691. text_name = "empty.txt"
  692. # Mock current_user
  693. with patch("services.file_service.current_user") as mock_current_user:
  694. mock_current_user.current_tenant_id = str(fake.uuid4())
  695. mock_current_user.id = str(fake.uuid4())
  696. upload_file = FileService.upload_text(text=text, text_name=text_name)
  697. assert upload_file is not None
  698. assert upload_file.size == 0
  699. def test_file_size_limits_edge_cases(self, db_session_with_containers, mock_external_service_dependencies):
  700. """
  701. Test file size limits with edge case values.
  702. """
  703. # Test exactly at limit
  704. for extension, limit_config in [
  705. ("jpg", dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT),
  706. ("mp4", dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT),
  707. ("mp3", dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT),
  708. ("pdf", dify_config.UPLOAD_FILE_SIZE_LIMIT),
  709. ]:
  710. file_size = limit_config * 1024 * 1024
  711. result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
  712. assert result is True
  713. # Test one byte over limit
  714. file_size = limit_config * 1024 * 1024 + 1
  715. result = FileService.is_file_size_within_limit(extension=extension, file_size=file_size)
  716. assert result is False
  717. def test_upload_file_with_source_url(self, db_session_with_containers, mock_external_service_dependencies):
  718. """
  719. Test file upload with source URL that gets overridden by signed URL.
  720. """
  721. fake = Faker()
  722. account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
  723. filename = "test.pdf"
  724. content = b"test content"
  725. mimetype = "application/pdf"
  726. source_url = "https://original-source.com/file.pdf"
  727. upload_file = FileService.upload_file(
  728. filename=filename,
  729. content=content,
  730. mimetype=mimetype,
  731. user=account,
  732. source_url=source_url,
  733. )
  734. # When source_url is provided, it should be preserved
  735. assert upload_file.source_url == source_url
  736. # The signed URL should only be set when source_url is empty
  737. # Let's test that scenario
  738. upload_file2 = FileService.upload_file(
  739. filename="test2.pdf",
  740. content=b"test content 2",
  741. mimetype="application/pdf",
  742. user=account,
  743. source_url="", # Empty source_url
  744. )
  745. # Should have the signed URL when source_url is empty
  746. assert upload_file2.source_url == "https://example.com/signed-url"