try_mock_using_exit_stack.py

from contextlib import ExitStack
import random
import unittest
from unittest.mock import patch


class TryMockUsingExitStackTest(unittest.TestCase):
    def test_mock_random_twice(self):
        with ExitStack() as stack:
            mock_randrange = stack.enter_context(
                patch('__main__.random.randrange')
            )
            mock_randint = stack.enter_context(
                patch('__main__.random.randint')
            )
            try_mock_using_exit_stack()
            mock_randint.assert_called_once_with(1, 2)
            mock_randrange.assert_called_once_with(3, 4)


def try_mock_using_exit_stack():
    random.randint(1, 2)
    random.randrange(3, 4)


unittest.main()