import unittest from pathlib import Path from PIL import Image from automation.engine import Matcher, Region, load_flow from automation.stamina import StaminaRunner ROOT = Path(__file__).resolve().parent.parent class StaminaDesktop: def __init__(self, images): self.images = images self.index = 0 self.clicks = [] def screenshot(self): return self.images[min(self.index, len(self.images) - 1)] def click(self, point, region, size): self.clicks.append(point) self.index += 1 class StaminaTests(unittest.TestCase): def setUp(self): self.flow = load_flow(ROOT / 'workflows/stamina_claim.json') self.matcher = Matcher(ROOT / 'assets') self.region = Region(0, 0, 562, 1002) def image(self, state): image = Image.new('RGB', (562, 1002), (42, 58, 73)) for spec in self.flow['states'][state]: template = Image.open(ROOT / 'assets' / spec['template']).convert('RGB') x, y, _, _ = spec['search'] image.paste(template, (x + 5, y + 5)) return image def runner(self, states): desktop = StaminaDesktop([self.image(state) for state in states]) return desktop, StaminaRunner(desktop, self.matcher, self.flow, lambda _: self.region, log=lambda _: None, timeout=0.8) def test_claims_stamina_and_returns_home(self): desktop, runner = self.runner(['home', 'territory', 'available', 'reward', 'empty', 'territory', 'home']) self.assertEqual(runner.run(), 'claimed') expected = ('territory', 'canteen', 'claim', 'dismiss_reward', 'close_canteen', 'territory_back') self.assertEqual(desktop.clicks, [self.flow['clicks'][name] for name in expected]) def test_empty_canteen_skips_claim_and_returns_home(self): desktop, runner = self.runner(['home', 'territory', 'empty', 'territory', 'home']) self.assertEqual(runner.run(), 'no_stamina') expected = ('territory', 'canteen', 'close_canteen', 'territory_back') self.assertEqual(desktop.clicks, [self.flow['clicks'][name] for name in expected]) if __name__ == '__main__': unittest.main()