diff --git a/botcity/core/bot.py b/botcity/core/bot.py index aa8728d..1a734ef 100644 --- a/botcity/core/bot.py +++ b/botcity/core/bot.py @@ -493,16 +493,20 @@ def find_same(item, items): return None haystack = self.get_screenshot() - eles = cv2find.locate_all_opencv( - element_path, - haystack_image=haystack, - region=region, - confidence=matching, - grayscale=grayscale, + # locate_all_opencv returns a generator, which is always truthy. + # Collect it so an empty result makes us look again. + eles = list( + cv2find.locate_all_opencv( + element_path, + haystack_image=haystack, + region=region, + confidence=matching, + grayscale=grayscale, + ) ) if not eles: continue - eles = deduplicate(list(eles)) + eles = deduplicate(eles) for ele in eles: if ele is not None: ele = self._fix_retina_element(ele) diff --git a/botcity/core/tests/test_find_all_wait.py b/botcity/core/tests/test_find_all_wait.py new file mode 100644 index 0000000..2d6cb96 --- /dev/null +++ b/botcity/core/tests/test_find_all_wait.py @@ -0,0 +1,32 @@ +import numpy +from PIL import Image + +from botcity.core import DesktopBot, cv2find + + +def test_find_all_keeps_looking_until_the_element_appears(tmp_path, monkeypatch): + rng = numpy.random.default_rng(0) + needle = rng.integers(0, 256, (10, 10, 3), dtype=numpy.uint8) + empty_screen = rng.integers(0, 256, (100, 200, 3), dtype=numpy.uint8) + loaded_screen = empty_screen.copy() + loaded_screen[20:30, 40:50] = needle + # The element shows up on the third screenshot, well within waiting_time. + screens = [empty_screen, empty_screen, loaded_screen] + + needle_path = tmp_path / "needle.png" + Image.fromarray(needle).save(needle_path) + bot = DesktopBot() + bot.add_image("needle", str(needle_path)) + monkeypatch.setattr(bot, "_fix_display_size", lambda: (200, 100)) + monkeypatch.setattr( + bot, + "get_screenshot", + lambda *a, **k: Image.fromarray( + screens.pop(0) if len(screens) > 1 else screens[0] + ), + ) + monkeypatch.setattr("botcity.core.bot.is_retina", lambda: False) + + found = list(bot.find_all("needle", matching=0.9, waiting_time=5000)) + + assert found == [cv2find.Box(40, 20, 10, 10)]