test_commons.py 2.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556
  1. import io
  2. import json
  3. import os
  4. import unittest
  5. from libs.commons import fitz
  6. from spark.s3_tools import get_s3_config, get_s3_client
  7. from libs.commons import join_path, json_dump_path, read_file, parse_bucket_key
  8. from loguru import logger
  9. test_pdf_dir_path = "s3://llm-pdf-text/unittest/pdf/"
  10. def get_test_pdf_json(book_name):
  11. json_path = join_path(json_dump_path, book_name + ".json")
  12. s3_config = get_s3_config(json_path)
  13. file_content = read_file(json_path, s3_config)
  14. json_str = file_content.decode('utf-8')
  15. json_object = json.loads(json_str)
  16. return json_object
  17. def read_test_file(book_name):
  18. test_pdf_path = join_path(test_pdf_dir_path, book_name + ".pdf")
  19. s3_config = get_s3_config(test_pdf_path)
  20. try:
  21. file_content = read_file(test_pdf_path, s3_config)
  22. return file_content
  23. except Exception as e:
  24. if "NoSuchKey" in str(e):
  25. logger.warning("File not found in test_pdf_path. Downloading from orig_s3_pdf_path.")
  26. try:
  27. json_object = get_test_pdf_json(book_name)
  28. orig_s3_pdf_path = json_object.get('file_location')
  29. s3_config = get_s3_config(orig_s3_pdf_path)
  30. file_content = read_file(orig_s3_pdf_path, s3_config)
  31. s3_client = get_s3_client(test_pdf_path)
  32. bucket_name, bucket_key = parse_bucket_key(test_pdf_path)
  33. file_obj = io.BytesIO(file_content)
  34. s3_client.upload_fileobj(file_obj, bucket_name, bucket_key)
  35. return file_content
  36. except Exception as e:
  37. logger.exception(e)
  38. else:
  39. logger.exception(e)
  40. def get_docs_from_test_pdf(book_name):
  41. file_content = read_test_file(book_name)
  42. return fitz.open("pdf", file_content)
  43. def get_test_json_data(directory_path, json_file_name):
  44. with open(os.path.join(directory_path, json_file_name), "r", encoding='utf-8') as f:
  45. test_data = json.load(f)
  46. return test_data