How to mock nested function calls of the same class

Question

In the post How to mock nested functions? the user mocks a chain of user calls. In my case I have a situation like the following:

class myClass:

    def __init__(self, my_type):
        self.type = my_type

    def output(self):
        self.create_figure()

    def create_figure():
        if self.type == 'A':
            do_A()
        elif self.type == 'B':
            do_B()

    def do_A():
        pass

    def do_B():
        pass

I am trying to write a unit test that tests whether the correct things were called. For example: I want to ensure when the user calls output on myClass('A'), that create_figure and do_A() are called.

I am currently testing with pytest as follows:

import myModule, pytest
from unittest.mock import patch

@pytest.fixture()
def default_my_class():
    return myModule.myClass(my_type='A')

@patch('myModule.myClass.do_A')
@patch('myModule.myClass.create_figure')
def test_chart(mock_create_figure, mock_do_a, default_my_class):
    default_my_class.output()
    mock_create_figure.assert_called_once()
    mock_do_a.assert_called_once()