iaai-parser/tests/test_self_heal.py

from __future__ import annotations

import json
import unittest
from unittest.mock import MagicMock, patch

from iaai_scraper.worker import self_heal


class TestSelfHeal(unittest.TestCase):
	def test_read_last_progress_ts_uses_global_key(self) -> None:
		redis_client = MagicMock()
		redis_client.get.return_value = "1776800000"

		ts = self_heal._read_last_progress_ts(redis_client)

		self.assertEqual(ts, 1776800000)
		redis_client.scan_iter.assert_not_called()

	def test_read_last_progress_ts_fallbacks_to_task_progress_keys(self) -> None:
		redis_client = MagicMock()
		redis_client.get.side_effect = lambda key: {
			self_heal.GLOBAL_PROGRESS_TS_KEY: None,
			"iaai:state:task_progress:a": '{"ts": 100}',
			"iaai:state:task_progress:b": '{"ts": 250}',
			"iaai:state:task_progress:c": '{"ts": 150}',
		}.get(key)
		redis_client.scan_iter.return_value = [
			"iaai:state:task_progress:a",
			"iaai:state:task_progress:b",
			"iaai:state:task_progress:c",
		]

		ts = self_heal._read_last_progress_ts(redis_client)

		self.assertEqual(ts, 250)

	def test_read_last_progress_ts_ignores_broken_payloads(self) -> None:
		redis_client = MagicMock()
		redis_client.get.side_effect = lambda key: {
			self_heal.GLOBAL_PROGRESS_TS_KEY: None,
			"iaai:state:task_progress:a": "{bad-json}",
			"iaai:state:task_progress:b": '{"foo": "bar"}',
		}.get(key)
		redis_client.scan_iter.return_value = [
			"iaai:state:task_progress:a",
			"iaai:state:task_progress:b",
		]

		ts = self_heal._read_last_progress_ts(redis_client)

		self.assertIsNone(ts)

	def test_active_recent_progress_prevents_db_idle_restart(self) -> None:
		redis_client = MagicMock()
		redis_client.scan_iter.return_value = ["iaai:state:task_progress:hourly"]
		redis_client.get.return_value = json.dumps(
			{
				"task_id": "hourly",
				"stage": "fast_listing_collected",
				"ts": 1000,
				"last_db_progress_ts": 1,
				"skipped_existing": 2417,
			}
		)

		active = self_heal._has_active_recent_progress(
			redis_client,
			now_ts=1010,
			stall_seconds=900,
		)

		self.assertTrue(active)

	@patch("iaai_scraper.worker.self_heal.time.time", return_value=5000)
	def test_has_inflight_work_deletes_stale_terminal_progress_without_work(self, _time_mock) -> None:
		redis_client = MagicMock()
		redis_client.llen.return_value = 0
		redis_client.get.side_effect = lambda key: {
			self_heal.SYNC_LISTING_LOCK_KEY: None,
			"iaai:state:task_progress:old": json.dumps(
				{
					"task_id": "old",
					"stage": "segment_done",
					"ts": 1000,
					"segment_full_scan_completed": True,
				}
			),
		}.get(key)
		redis_client.scan_iter.return_value = ["iaai:state:task_progress:old"]

		has_inflight, flags = self_heal._has_inflight_work(redis_client, self_heal.IAAI_SYNC_QUEUE)

		self.assertFalse(has_inflight)
		self.assertEqual(flags["has_task_progress"], 0)
		redis_client.delete.assert_called_once_with("iaai:state:task_progress:old")

	@patch("iaai_scraper.worker.self_heal.time.time", return_value=1010)
	def test_has_inflight_work_keeps_recent_non_terminal_progress(self, _time_mock) -> None:
		redis_client = MagicMock()
		redis_client.llen.return_value = 0
		redis_client.get.side_effect = lambda key: {
			self_heal.SYNC_LISTING_LOCK_KEY: None,
			"iaai:state:task_progress:active": json.dumps(
				{
					"task_id": "active",
					"stage": "fast_listing_collected",
					"ts": 1000,
				}
			),
		}.get(key)
		redis_client.scan_iter.return_value = ["iaai:state:task_progress:active"]

		has_inflight, flags = self_heal._has_inflight_work(redis_client, self_heal.IAAI_SYNC_QUEUE)

		self.assertTrue(has_inflight)
		self.assertEqual(flags["has_task_progress"], 1)
		redis_client.delete.assert_not_called()

	@patch("iaai_scraper.worker.self_heal.time.time", return_value=5000)
	def test_has_inflight_work_deletes_stale_non_terminal_progress_without_work(self, _time_mock) -> None:
		redis_client = MagicMock()
		redis_client.llen.return_value = 0
		redis_client.get.side_effect = lambda key: {
			self_heal.SYNC_LISTING_LOCK_KEY: None,
			"iaai:state:task_progress:stale": json.dumps(
				{
					"task_id": "stale",
					"stage": "segment_started",
					"ts": 1000,
					"segment_index": 11,
				}
			),
		}.get(key)
		redis_client.scan_iter.return_value = ["iaai:state:task_progress:stale"]

		has_inflight, flags = self_heal._has_inflight_work(redis_client, self_heal.IAAI_SYNC_QUEUE)

		self.assertFalse(has_inflight)
		self.assertEqual(flags["has_task_progress"], 0)
		self.assertEqual(flags["stale_progress_deleted"], 1)
		redis_client.delete.assert_called_once_with("iaai:state:task_progress:stale")

	@patch("iaai_scraper.worker.self_heal.time.time", return_value=5000)
	def test_has_inflight_work_keeps_stale_progress_when_lock_exists(self, _time_mock) -> None:
		redis_client = MagicMock()
		redis_client.llen.return_value = 0
		redis_client.get.side_effect = lambda key: {
			self_heal.SYNC_LISTING_LOCK_KEY: "owner",
			"iaai:state:task_progress:stale": json.dumps(
				{
					"task_id": "stale",
					"stage": "segment_started",
					"ts": 1000,
				}
			),
		}.get(key)
		redis_client.scan_iter.return_value = ["iaai:state:task_progress:stale"]

		has_inflight, flags = self_heal._has_inflight_work(redis_client, self_heal.IAAI_SYNC_QUEUE)

		self.assertTrue(has_inflight)
		self.assertEqual(flags["has_lock"], 1)
		self.assertEqual(flags["has_task_progress"], 1)
		redis_client.delete.assert_not_called()

	@patch("iaai_scraper.worker.self_heal.time.sleep", return_value=None)
	@patch("iaai_scraper.worker.self_heal.os.kill")
	@patch("builtins.open")
	def test_kill_worker_process_sends_term_and_kill(self, open_mock, kill_mock, _sleep_mock) -> None:
		open_mock.return_value.__enter__.return_value.read.return_value = "123"
		# SIGTERM -> process alive check (pid,0) -> SIGKILL
		kill_mock.side_effect = [None, None, None]

		self_heal._kill_worker_process()

		self.assertEqual(kill_mock.call_args_list[0].args[0], 123)
		self.assertEqual(kill_mock.call_args_list[1].args, (123, 0))
		self.assertEqual(kill_mock.call_args_list[2].args[0], 123)

	@patch("iaai_scraper.worker.self_heal.time.sleep", return_value=None)
	@patch("iaai_scraper.worker.self_heal.os.kill")
	@patch("builtins.open")
	def test_kill_worker_process_skips_sigkill_when_already_exited(self, open_mock, kill_mock, _sleep_mock) -> None:
		open_mock.return_value.__enter__.return_value.read.return_value = "123"
		kill_mock.side_effect = [None, ProcessLookupError()]

		self_heal._kill_worker_process()

		# Только SIGTERM и проверка существования процесса.
		self.assertEqual(len(kill_mock.call_args_list), 2)


if __name__ == "__main__":
	unittest.main()