How to set an object reference field in the Inspector

What I want is to be able to call a method on an object (doesn't need to be MonoBehaviour), where the specific implementation of that method is chosen in the inspector. I will paste an example of what I'm trying to do, then describe in better detail how I would like it to work.

The Base Class (Could also be an Interface, Abstract class... whatever works.)

 [System.Serializable]
 public class TestBase
 {
     public virtual string RetStr()
     {
         return "Base Class";
     }
 }

The inherited class/Implementer:

 [System.Serializable]
 public class TestClass : TestBase
 {
     public override string RetStr()
     {
         return "Inherited Class";
     }
 }

The MonoBehaviour, with the public field I wish to set:

 using UnityEngine;
 
 public class Assignee : MonoBehaviour {
     public TestBase testObj;
 
     private void Awake()
     {
         if (testObj == null)
             Debug.Log("testObj is null");
         else
             Debug.Log("testObj.RetStr() = " + testObj.RetStr());
     }
 }

And the custom editor (that, to keep it simple, Assigns either the base or the inherited):

 using UnityEngine;
 using UnityEditor;
 
 [CustomEditor(typeof(Assignee))]
 public class AssigneeEditor : Editor
 {
     Assignee assignee;
     bool setTest;
 
     void OnEnable()
     {
         assignee = target as Assignee;
         setTest = assignee.testObj.GetType() != typeof(TestBase);
     }
 
     public override void OnInspectorGUI()
     {
         setTest = EditorGUILayout.Toggle(new GUIContent("Set testObj"), setTest);
 
         if (setTest)
             assignee.testObj = new TestClass();
         else
             assignee.testObj = new TestBase();
 
         serializedObject.ApplyModifiedProperties();
     }
 }

The issue is that, while I can change the field in play mode, entering/exiting play mode resets it, also that this whole method is weird. This is probably due to it not being serialized.

The goal is for the workflow to be like: Create new C# script -> Change it to inherit from/implement TestBase -> On a MonoBehaviour component that has a TestBase field, choose a specific one somehow (for now, just selecting between two is sufficient to prove it can be done).