test_simulation.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. import asyncio
  2. import os
  3. import sys
  4. import shutil
  5. import json
  6. from unittest.mock import MagicMock, AsyncMock, patch
  7. # Add project root to path
  8. sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
  9. from app.database import Base, engine, SessionLocal
  10. from app.services.webhook_service import WebhookService
  11. from app.models import Project, DataVersion, DataFile
  12. # Setup
  13. def setup_db():
  14. Base.metadata.drop_all(bind=engine)
  15. Base.metadata.create_all(bind=engine)
  16. def teardown_storage():
  17. if os.path.exists("test_storage"):
  18. shutil.rmtree("test_storage")
  19. # Mocks
  20. manifest_yaml = """
  21. project_name: "test_project"
  22. stage: "test_stage"
  23. outputs:
  24. - path: "results/data.csv"
  25. - path: "images/"
  26. pattern: "*.png"
  27. """
  28. mock_tree_response = {
  29. "tree": [
  30. {"path": "results/data.csv", "type": "blob", "sha": "sha123", "mode": "100644"},
  31. {"path": "images/plot.png", "type": "blob", "sha": "sha456", "mode": "100644"},
  32. {"path": "images/ignore.txt", "type": "blob", "sha": "sha789", "mode": "100644"},
  33. {"path": "README.md", "type": "blob", "sha": "sha000", "mode": "100644"}
  34. ]
  35. }
  36. async def run_test():
  37. print("Setting up test environment...")
  38. setup_db()
  39. teardown_storage()
  40. # Override settings for storage root
  41. with patch("app.config.settings.STORAGE_ROOT", "test_storage"):
  42. db = SessionLocal()
  43. service = WebhookService(db)
  44. # Mock GogsClient
  45. service.gogs.get_manifest = AsyncMock(return_value=manifest_yaml)
  46. service.gogs.get_recursive_tree = AsyncMock(return_value=mock_tree_response)
  47. async def mock_get_file(owner, repo, commit, path):
  48. if path == "results/data.csv":
  49. return b"csv_data"
  50. if path == "images/plot.png":
  51. return b"png_data"
  52. return b"other_data"
  53. service.gogs.get_file_content = AsyncMock(side_effect=mock_get_file)
  54. # Mock Payload
  55. payload = {
  56. "ref": "refs/heads/master",
  57. "after": "commit_sha_abc",
  58. "repository": {
  59. "name": "my-repo",
  60. "owner": {"username": "my-user"}
  61. },
  62. "pusher": {"username": "test-author"}
  63. }
  64. print("Processing webhook...")
  65. await service.process_webhook(payload)
  66. # Verification
  67. print("Verifying results...")
  68. # Check Project
  69. project = db.query(Project).filter_by(project_name="test_project").first()
  70. assert project is not None
  71. print("[PASS] Project created")
  72. # Check Version
  73. version = db.query(DataVersion).filter_by(commit_id="commit_sha_abc").first()
  74. assert version is not None
  75. assert version.stage == "test_stage"
  76. print("[PASS] Version created")
  77. # Check Files
  78. files = db.query(DataFile).filter_by(version_id=version.id).all()
  79. file_paths = [f.relative_path for f in files]
  80. print(f"Stored files: {file_paths}")
  81. assert "results/data.csv" in file_paths
  82. assert "images/plot.png" in file_paths
  83. # ignore.txt should check against pattern if pattern logic is strict?
  84. # My pattern logic was: if dir, check pattern.
  85. # manifest: images/ pattern: *.png. So ignore.txt should NOT be there.
  86. # But wait, logic in webhook_service: file_path.startswith(path_pattern) ... fnmatch(rel_name, match_pattern)
  87. # ignore.txt -> rel_name: ignore.txt. *.png matches? No. Good.
  88. if "images/ignore.txt" in file_paths:
  89. print("[FAIL] 'images/ignore.txt' should not be included")
  90. else:
  91. print("[PASS] Pattern filtering worked")
  92. # Check Physical Files
  93. for f in files:
  94. if not os.path.exists(f.storage_path):
  95. print(f"[FAIL] {f.storage_path} does not exist")
  96. else:
  97. print(f"[PASS] File {f.relative_path} stored at {f.storage_path}")
  98. # Test Deduplication
  99. print("\nTesting Deduplication...")
  100. # New payload, same file content (same SHA)
  101. payload2 = payload.copy()
  102. payload2["after"] = "commit_sha_def" # New commit
  103. # We need to simulate that this new commit has the SAME tree for these files
  104. # service.gogs.get_recursive_tree is already mocked to return the same tree SHAs
  105. # Reset get_file_content mock to track calls
  106. service.gogs.get_file_content.reset_mock()
  107. await service.process_webhook(payload2)
  108. version2 = db.query(DataVersion).filter_by(commit_id="commit_sha_def").first()
  109. assert version2 is not None
  110. print("[PASS] Version 2 created")
  111. # Check if get_file_content was called. It should NOT be called because SHAs are same and files exist.
  112. if service.gogs.get_file_content.called:
  113. print("[FAIL] Deduplication failed: files were downloaded again")
  114. print(service.gogs.get_file_content.mock_calls)
  115. else:
  116. print("[PASS] Deduplication worked: download skipped")
  117. db.close()
  118. teardown_storage()
  119. print("\nAll tests passed!")
  120. if __name__ == "__main__":
  121. asyncio.run(run_test())