213 lines
8.4 KiB
Python
213 lines
8.4 KiB
Python
import hashlib
|
||
import json
|
||
import os
|
||
import signal
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import textwrap
|
||
import unittest
|
||
|
||
from _common import SRC
|
||
|
||
if SRC not in sys.path:
|
||
sys.path.insert(0, SRC)
|
||
import run # noqa: E402
|
||
|
||
# Поддельные агенты: маленькие скрипты, которые ведут себя как agent.py по контракту (коды выхода, update.json).
|
||
STATE_HELPERS = textwrap.dedent('''
|
||
import json, os, sys
|
||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||
STATE = os.path.join(HERE, "update.json")
|
||
def state():
|
||
try:
|
||
return json.load(open(STATE, encoding="utf-8"))
|
||
except Exception:
|
||
return None
|
||
def write(**kw):
|
||
json.dump(kw, open(STATE, "w", encoding="utf-8"))
|
||
def mark(name):
|
||
open(os.path.join(HERE, "runs.log"), "a").write(name + "\\n")
|
||
''')
|
||
|
||
|
||
def fake(name, body):
|
||
return STATE_HELPERS + "mark(%r)\n" % name + textwrap.dedent(body)
|
||
|
||
|
||
# v5 «нашёл» новую версию: кладёт agent.py.new из next.py, пишет staged и выходит с 75; после отката/отказа — выход 0.
|
||
V5_UPDATES = fake("v5", '''
|
||
import hashlib
|
||
st = state()
|
||
if st is None:
|
||
data = open(os.path.join(HERE, "next.py"), "rb").read()
|
||
open(os.path.join(HERE, "agent.py.new"), "wb").write(data)
|
||
write(state="staged", from_version=5, to_version=6, sha256=hashlib.sha256(data).hexdigest())
|
||
sys.exit(75)
|
||
sys.exit(0)
|
||
''')
|
||
# v5 с подменённым после проверки файлом: sha в staged не сходится с agent.py.new.
|
||
V5_BAD_STAGE = fake("v5", '''
|
||
st = state()
|
||
if st is None:
|
||
open(os.path.join(HERE, "agent.py.new"), "w").write("мусор")
|
||
write(state="staged", from_version=5, to_version=6, sha256="чужой")
|
||
sys.exit(75)
|
||
sys.exit(0)
|
||
''')
|
||
V6_OK_THEN_STOP = fake("v6", '''
|
||
st = state()
|
||
write(**dict(st, state="ok"))
|
||
sys.exit(0)
|
||
''')
|
||
V6_CRASH = fake("v6", 'sys.exit(1)\n')
|
||
V6_HANG = fake("v6", 'import time\ntime.sleep(60)\n')
|
||
|
||
|
||
class LauncherTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.dir = self.tmp.name
|
||
self.l = run.Launcher(self.dir, ["--journal", "j.jsonl"])
|
||
self.pauses = []
|
||
self.l.pause = lambda s: (self.pauses.append(s), True)[1]
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def put(self, name, text):
|
||
with open(os.path.join(self.dir, name), "w", encoding="utf-8") as f:
|
||
f.write(text)
|
||
|
||
def read(self, name):
|
||
with open(os.path.join(self.dir, name), encoding="utf-8") as f:
|
||
return f.read()
|
||
|
||
def state(self):
|
||
return json.loads(self.read("update.json"))
|
||
|
||
def runs(self):
|
||
return self.read("runs.log").split()
|
||
|
||
def next_version(self, text):
|
||
self.put("next.py", text)
|
||
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||
|
||
def stage(self, new_text, to_version=6):
|
||
self.put("agent.py.new", new_text)
|
||
sha = hashlib.sha256(new_text.encode("utf-8")).hexdigest()
|
||
self.put("update.json", json.dumps({"state": "staged", "from_version": 5, "to_version": to_version, "sha256": sha}))
|
||
return sha
|
||
|
||
def test_update_applied_and_confirmed(self):
|
||
self.put("agent.py", V5_UPDATES)
|
||
sha = self.next_version(V6_OK_THEN_STOP)
|
||
self.assertEqual(self.l.run(), 0)
|
||
self.assertEqual(self.read("agent.py"), V6_OK_THEN_STOP)
|
||
self.assertEqual(self.read("agent.py.prev"), V5_UPDATES)
|
||
self.assertFalse(os.path.exists(os.path.join(self.dir, "agent.py.new")))
|
||
st = self.state()
|
||
self.assertEqual((st["state"], st["from_version"], st["to_version"], st["sha256"]), ("ok", 5, 6, sha))
|
||
self.assertEqual(self.runs(), ["v5", "v6"])
|
||
|
||
def test_crash_in_trial_rolls_back(self):
|
||
self.put("agent.py", V5_UPDATES)
|
||
sha = self.next_version(V6_CRASH)
|
||
self.assertEqual(self.l.run(), 0)
|
||
self.assertEqual(self.read("agent.py"), V5_UPDATES)
|
||
st = self.state()
|
||
self.assertEqual((st["state"], st["to_version"], st["sha256"]), ("rolled_back", 6, sha))
|
||
self.assertRegex(st["error"], r"^код 1 через \d+ с$")
|
||
self.assertEqual(self.runs(), ["v5", "v6", "v5"])
|
||
self.assertEqual(self.pauses, []) # откат без паузы
|
||
|
||
def test_hang_in_trial_rolls_back(self):
|
||
self.l.trial_seconds = 1
|
||
self.l.kill_grace = 5
|
||
self.put("agent.py", V5_UPDATES)
|
||
self.next_version(V6_HANG)
|
||
self.assertEqual(self.l.run(), 0)
|
||
self.assertEqual(self.read("agent.py"), V5_UPDATES)
|
||
self.assertEqual(self.state()["error"], "не вышла на связь за 10 мин")
|
||
|
||
def test_crash_after_ok_restarts_with_backoff(self):
|
||
self.put("update.json", json.dumps({"state": "ok", "from_version": 5, "to_version": 6, "sha256": "s"}))
|
||
self.put("agent.py", fake("v6", '''
|
||
n = len(open(os.path.join(HERE, "runs.log")).read().split())
|
||
sys.exit(1 if n <= 2 else 0)
|
||
'''))
|
||
self.assertEqual(self.l.run(), 0)
|
||
self.assertEqual(self.pauses, [5, 10])
|
||
self.assertEqual(self.state()["state"], "ok") # отката нет
|
||
|
||
def test_backoff_caps_and_resets(self):
|
||
self.assertEqual([self.l.next_backoff(10) for _ in range(8)], [5, 10, 20, 40, 80, 160, 300, 300])
|
||
self.assertEqual(self.l.next_backoff(600), 5) # проработал 10 минут — с начала
|
||
|
||
def test_75_without_staged_rejected(self):
|
||
self.put("agent.py", V5_BAD_STAGE)
|
||
self.assertEqual(self.l.run(), 0)
|
||
self.assertEqual(self.read("agent.py"), V5_BAD_STAGE)
|
||
self.assertEqual(self.runs(), ["v5", "v5"])
|
||
self.assertFalse(os.path.exists(os.path.join(self.dir, "agent.py.new")))
|
||
st = self.state()
|
||
self.assertEqual((st["state"], st["error"]), ("rejected", "файл изменился после проверки"))
|
||
|
||
def test_staged_at_start_is_applied(self):
|
||
self.put("agent.py", fake("v5", "sys.exit(0)\n"))
|
||
self.stage(V6_OK_THEN_STOP)
|
||
self.assertEqual(self.l.run(), 0)
|
||
self.assertEqual(self.runs(), ["v6"])
|
||
self.assertEqual(self.state()["state"], "ok")
|
||
|
||
def test_stray_new_removed_at_start(self):
|
||
self.put("agent.py", fake("v5", "sys.exit(0)\n"))
|
||
self.put("agent.py.new", "остаток")
|
||
self.assertEqual(self.l.run(), 0)
|
||
self.assertFalse(os.path.exists(os.path.join(self.dir, "agent.py.new")))
|
||
|
||
def test_missing_agent_restored_from_prev(self):
|
||
self.put("agent.py.prev", fake("v5", "sys.exit(0)\n"))
|
||
self.assertEqual(self.l.run(), 0)
|
||
self.assertEqual(self.runs(), ["v5"])
|
||
|
||
def test_garbage_update_json_is_no_state(self):
|
||
self.put("update.json", '{"state": "tri')
|
||
self.put("agent.py", fake("v5", "sys.exit(0)\n"))
|
||
self.assertEqual(self.l.run(), 0)
|
||
|
||
def test_env_and_args_passed(self):
|
||
self.put("agent.py", fake("v5", '''
|
||
open(os.path.join(HERE, "env.txt"), "w").write(os.environ.get("JIRA_AGENT_LAUNCHER", "") + " " + " ".join(sys.argv[1:]))
|
||
sys.exit(130)
|
||
'''))
|
||
self.assertEqual(self.l.run(), 130)
|
||
self.assertEqual(self.read("env.txt"), "1 run --journal j.jsonl")
|
||
|
||
def test_sigint_without_child_stops(self):
|
||
self.l.on_sigint(signal.SIGINT, None)
|
||
self.assertTrue(self.l.stopping)
|
||
|
||
def test_sigint_with_child_ignored(self):
|
||
self.l.child = object()
|
||
self.l.on_sigint(signal.SIGINT, None)
|
||
self.assertFalse(self.l.stopping)
|
||
|
||
def test_second_launcher_refused(self):
|
||
with open(os.path.join(SRC, "run.py"), encoding="utf-8") as f:
|
||
self.put("run.py", f.read())
|
||
self.put("agent.py", fake("v5", "sys.exit(0)\n"))
|
||
lock = run.acquire_lock(self.dir)
|
||
try:
|
||
p = subprocess.run([sys.executable, os.path.join(self.dir, "run.py")], stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT, timeout=30)
|
||
finally:
|
||
lock.close()
|
||
self.assertEqual(p.returncode, 1)
|
||
self.assertIn("запускатель уже работает", p.stdout.decode("utf-8"))
|
||
self.assertFalse(os.path.exists(os.path.join(self.dir, "runs.log")))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|