How to make FactoryBot return the right STI sub class?

I'm making a big change in my system, so I changed one of my main tables into a STI, and create subclasses to implement the specific behavior.

class MainProcess < ApplicationRecord
end

class ProcessA < MainProcess
end

class ProcessB < MainProcess
end

In the application code, if I run MainProcess.new(type: 'ProcessA') it will return a ProcessA as I want. But in the Rspec tests when I run FactoryBot::create(:main_process, type: 'ProcessA') it is returning a MainProcess and breaking my tests.

My factor is something like this

FactoryBot.define do
  factory :main_process do
    foo { 'bar' }
  end

  factory :process_a, parent: :main_process, class: 'ProcessA' do
  end

  factory :process_b, parent: :main_process, class: 'ProcessB' do
  end
end

Is there some way to make FactoryBot have the same behavior of normal program?

I found the solution

FactoryBot.define do
  factory :main_process do
    initialize_with do
      klass = type.constantize
      klass.new(attributes)  
    end
  end
  ...
end

The answer was founded here http://indigolain.hatenablog.com/entry/defining-factory-for-sti-defined-model (in japanese)

Edit #1:

⚠⚠⚠ Important ⚠⚠⚠

As mentioned here initialize_with is part of a private FactoryBot API.

According to the documentation:

This method is part of a private API. You should avoid using this method if possible, as it may be removed or be changed in the future.

So avoid to use if you can. (although I didn't find any other way to achieve this result without use it)

Edit #2

Besides the warning in the gem documentation (described above), the GETTING_STARTED.md actually suggest you use it

If you want to use factory_bot to construct an object where some attributes are passed to initialize or if you want to do something other than simply calling new on your build class, you can override the default behavior by defining initialize_with on your factory