65 lines
2.8 KiB
Python
65 lines
2.8 KiB
Python
"""Fake `claude -p` for the CLI tests. $FAKE_CLAUDE_SCENARIO: ok (default), miss, api_error,
|
|
noresult, two_models, no_cost, sleep. $FAKE_CLAUDE_LOG: one JSON line per call (argv, pid)."""
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import uuid
|
|
|
|
TRAIN_MU = {"inputTokens": 10, "outputTokens": 20, "cacheReadInputTokens": 0,
|
|
"cacheCreationInputTokens": 1000}
|
|
TRAIN_COST, CALL_COST = 0.05, 0.02
|
|
|
|
|
|
def opt(name):
|
|
return sys.argv[sys.argv.index(name) + 1] if name in sys.argv else None
|
|
|
|
|
|
def emit(ev):
|
|
print(json.dumps(ev), flush=True)
|
|
|
|
|
|
def main():
|
|
prompt = sys.stdin.buffer.read().decode("utf-8")
|
|
scen = os.environ.get("FAKE_CLAUDE_SCENARIO", "ok")
|
|
if os.environ.get("FAKE_CLAUDE_LOG"):
|
|
with open(os.environ["FAKE_CLAUDE_LOG"], "a", encoding="utf-8") as f:
|
|
f.write(json.dumps({"argv": sys.argv, "pid": os.getpid(), "prompt_chars": len(prompt)}) + "\n")
|
|
resume, model = opt("--resume"), opt("--model")
|
|
# a plain resume (no fork, no --session-id) gets a new id, as `-p` may give one
|
|
sid = opt("--session-id") or ("resumed-" + uuid.uuid4().hex if resume else str(uuid.uuid4()))
|
|
read = 1000 if resume and scen != "miss" else 0
|
|
created = 100 if resume else 1000
|
|
usage = {"input_tokens": 10, "output_tokens": 20, "cache_creation_input_tokens": created,
|
|
"cache_read_input_tokens": read,
|
|
"cache_creation": {"ephemeral_1h_input_tokens": created, "ephemeral_5m_input_tokens": 0}}
|
|
emit({"type": "system", "subtype": "init", "session_id": sid, "model": model})
|
|
emit({"type": "assistant", "session_id": sid, "message": {"id": "m1", "usage": usage}})
|
|
if scen == "sleep":
|
|
time.sleep(60)
|
|
if scen == "noresult":
|
|
sys.exit(1)
|
|
mu = {"inputTokens": 10, "outputTokens": 20, "cacheReadInputTokens": read,
|
|
"cacheCreationInputTokens": created}
|
|
cost = CALL_COST if resume else TRAIN_COST
|
|
if resume: # the session's totals include the base's training call
|
|
mu = {k: v + TRAIN_MU[k] for k, v in mu.items()}
|
|
cost += TRAIN_COST
|
|
res = {"type": "result", "subtype": "success", "is_error": False, "session_id": sid,
|
|
"result": "done", "num_turns": 1, "usage": usage, "modelUsage": {model: mu},
|
|
"total_cost_usd": cost}
|
|
if scen == "two_models":
|
|
res["modelUsage"]["claude-haiku-5-5"] = {"inputTokens": 5, "outputTokens": 1,
|
|
"cacheReadInputTokens": 0, "cacheCreationInputTokens": 0}
|
|
if scen == "no_cost":
|
|
del res["total_cost_usd"]
|
|
if scen == "api_error":
|
|
res.update(subtype="error_during_execution", is_error=True, api_error_status=429,
|
|
result="API Error: 429 rate limit")
|
|
emit(res)
|
|
sys.exit(1 if scen == "api_error" else 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|