badfish/tests/test_async_loop.py at 9f3523cab31a75e9bede6c7f058c7878d7e09675 · redhat-performance/badfish · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import unittest
from unittest.mock import patch, MagicMock
from badfish.main import main


class TestAsyncioFix(unittest.TestCase):
    @patch("badfish.main.execute_badfish")
    @patch("badfish.main.BadfishLogger")
    @patch("badfish.main.parse_arguments")
    @patch("asyncio.set_event_loop")
    @patch("asyncio.new_event_loop")
    @patch("asyncio.get_event_loop")
    def test_main_handles_no_event_loop(
        self, mock_get_loop, mock_new_loop, mock_set_loop, mock_parse_args, mock_logger, mock_execute
    ):
        mock_get_loop.side_effect = RuntimeError("No event loop")

        mock_loop_instance = MagicMock()
        mock_new_loop.return_value = mock_loop_instance
        mock_loop_instance.run_until_complete.return_value = ("localhost", True)

        mock_parse_args.return_value = {
            "verbose": False,
            "host": "localhost",
            "delta": None,
            "firmware_inventory": None,
            "host_list": None,
            "log": None,
            "output": None,
        }

        main()

        mock_get_loop.assert_called_once()
        mock_new_loop.assert_called_once()
        mock_set_loop.assert_called_once_with(mock_loop_instance)
        mock_loop_instance.run_until_complete.assert_called()

    @patch("badfish.main.execute_badfish")
    @patch("badfish.main.BadfishLogger")
    @patch("badfish.main.parse_arguments")
    @patch("asyncio.set_event_loop")
    @patch("asyncio.new_event_loop")
    @patch("asyncio.get_event_loop")
    def test_main_uses_existing_loop(
        self, mock_get_loop, mock_new_loop, mock_set_loop, mock_parse_args, mock_logger, mock_execute
    ):
        existing_loop = MagicMock()
        mock_get_loop.return_value = existing_loop
        mock_get_loop.side_effect = None
        existing_loop.run_until_complete.return_value = ("localhost", True)

        mock_parse_args.return_value = {
            "verbose": False,
            "host": "localhost",
            "delta": None,
            "firmware_inventory": None,
            "host_list": None,
            "log": None,
            "output": None,
        }

        main()

        mock_get_loop.assert_called_once()
        mock_new_loop.assert_not_called()
        mock_set_loop.assert_not_called()
        existing_loop.run_until_complete.assert_called()