test_normalizer.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179
  1. import json
  2. from drill.normalizer import (
  3. collect_new_logs,
  4. filter_codex_logs_by_cwd,
  5. normalize_claude_logs,
  6. normalize_codex_logs,
  7. normalize_gemini_logs,
  8. snapshot_log_dir,
  9. )
  10. class TestSnapshotAndCollect:
  11. def test_snapshot_and_collect_new_files(self, tmp_path):
  12. log_dir = tmp_path / "logs"
  13. log_dir.mkdir()
  14. (log_dir / "old.jsonl").write_text('{"old": true}\n')
  15. snapshot = snapshot_log_dir(log_dir)
  16. (log_dir / "new.jsonl").write_text('{"new": true}\n')
  17. new_files = collect_new_logs(log_dir, snapshot)
  18. assert len(new_files) == 1
  19. assert new_files[0].name == "new.jsonl"
  20. def test_empty_dir_returns_empty(self, tmp_path):
  21. log_dir = tmp_path / "logs"
  22. log_dir.mkdir()
  23. snapshot = snapshot_log_dir(log_dir)
  24. new_files = collect_new_logs(log_dir, snapshot)
  25. assert new_files == []
  26. class TestNormalizeClaudeLogs:
  27. def test_normalizes_tool_use(self):
  28. lines = [
  29. json.dumps(
  30. {"type": "tool_use", "name": "EnterWorktree", "input": {"branch": "add-login"}}
  31. ),
  32. json.dumps({"type": "tool_use", "name": "Bash", "input": {"command": "git status"}}),
  33. json.dumps({"type": "text", "text": "I'll create a worktree"}),
  34. ]
  35. normalized = normalize_claude_logs("\n".join(lines))
  36. assert len(normalized) == 2
  37. assert normalized[0]["tool"] == "EnterWorktree"
  38. assert normalized[0]["source"] == "native"
  39. assert normalized[1]["tool"] == "Bash"
  40. assert normalized[1]["source"] == "shell"
  41. class TestNormalizeCodexLogs:
  42. def test_normalizes_local_shell_call(self):
  43. lines = [
  44. json.dumps(
  45. {
  46. "type": "response_item",
  47. "item": {
  48. "type": "local_shell_call",
  49. "action": {"command": ["git", "worktree", "add", "feature"]},
  50. "status": "completed",
  51. },
  52. }
  53. ),
  54. json.dumps(
  55. {
  56. "type": "response_item",
  57. "item": {"type": "message", "content": [{"text": "Creating worktree"}]},
  58. }
  59. ),
  60. ]
  61. normalized = normalize_codex_logs("\n".join(lines))
  62. assert len(normalized) == 1
  63. assert normalized[0]["tool"] == "Bash"
  64. assert "git worktree add" in normalized[0]["args"]["command"]
  65. assert normalized[0]["source"] == "shell"
  66. def test_filter_by_cwd_keeps_matching_drops_others(self, tmp_path):
  67. target = "/private/tmp/drill-target"
  68. match = tmp_path / "match.jsonl"
  69. match.write_text(
  70. json.dumps(
  71. {
  72. "type": "session_meta",
  73. "payload": {"id": "abc", "cwd": target},
  74. }
  75. )
  76. + "\n"
  77. )
  78. other = tmp_path / "other.jsonl"
  79. other.write_text(
  80. json.dumps(
  81. {
  82. "type": "session_meta",
  83. "payload": {"id": "def", "cwd": "/private/tmp/drill-other"},
  84. }
  85. )
  86. + "\n"
  87. )
  88. no_meta = tmp_path / "no-meta.jsonl"
  89. no_meta.write_text(json.dumps({"type": "response_item", "payload": {}}) + "\n")
  90. empty = tmp_path / "empty.jsonl"
  91. empty.write_text("")
  92. kept = filter_codex_logs_by_cwd([match, other, no_meta, empty], target)
  93. assert kept == [match]
  94. def test_normalizes_function_call_with_payload(self):
  95. """Test the actual codex rollout format using payload instead of item."""
  96. lines = [
  97. json.dumps(
  98. {
  99. "type": "response_item",
  100. "payload": {
  101. "type": "function_call",
  102. "name": "exec_command",
  103. "arguments": '{"cmd":"git worktree add .worktrees/feature",'
  104. '"workdir":"/tmp/test"}',
  105. "call_id": "call_123",
  106. },
  107. }
  108. ),
  109. json.dumps(
  110. {
  111. "type": "response_item",
  112. "payload": {
  113. "type": "function_call",
  114. "name": "apply_patch",
  115. "arguments": '{"patch":"--- a/file\\n+++ b/file"}',
  116. "call_id": "call_456",
  117. },
  118. }
  119. ),
  120. ]
  121. normalized = normalize_codex_logs("\n".join(lines))
  122. assert len(normalized) == 2
  123. assert normalized[0]["tool"] == "Bash"
  124. assert "git worktree add" in normalized[0]["args"]["command"]
  125. assert normalized[0]["source"] == "shell"
  126. assert normalized[1]["tool"] == "Edit"
  127. assert normalized[1]["source"] == "native"
  128. class TestNormalizeGeminiLogs:
  129. def test_normalizes_jsonl_tool_calls(self):
  130. lines = [
  131. json.dumps({"kind": "main"}),
  132. json.dumps(
  133. {
  134. "type": "gemini",
  135. "content": "Reading file",
  136. "toolCalls": [
  137. {
  138. "id": "read_file_1",
  139. "name": "read_file",
  140. "args": {"file_path": "GEMINI.md"},
  141. "status": "success",
  142. }
  143. ],
  144. }
  145. ),
  146. json.dumps(
  147. {
  148. "type": "gemini",
  149. "content": "Running command",
  150. "toolCalls": [
  151. {
  152. "id": "shell_1",
  153. "name": "run_shell_command",
  154. "args": {"command": "git status"},
  155. "status": "success",
  156. }
  157. ],
  158. }
  159. ),
  160. ]
  161. normalized = normalize_gemini_logs("\n".join(lines))
  162. assert normalized == [
  163. {"tool": "Read", "args": {"file_path": "GEMINI.md"}, "source": "native"},
  164. {"tool": "Bash", "args": {"command": "git status"}, "source": "shell"},
  165. ]