From 4280c54dd4a92f1cf8113c8be434e7ba108cdf69 Mon Sep 17 00:00:00 2001
From: Benjamin Michaelis <git@relay.benjamin.michaelis.net>
Date: Fri, 2 Feb 2024 14:55:03 -0800
Subject: [PATCH] Change return type

---
 Moq.AutoMock.Tests/DescribeWithSelfMock.cs | 13 +++++++++++--
 Moq.AutoMock/AutoMocker.cs                 |  6 +++---
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/Moq.AutoMock.Tests/DescribeWithSelfMock.cs b/Moq.AutoMock.Tests/DescribeWithSelfMock.cs
index ce4d5e5e..b6113378 100644
--- a/Moq.AutoMock.Tests/DescribeWithSelfMock.cs
+++ b/Moq.AutoMock.Tests/DescribeWithSelfMock.cs
@@ -12,7 +12,7 @@ public void It_can_register_a_custom_default_value_provider_for_a_self_mock_with
 
         var mock = mocker.WithSelfMock<IService2, Service2>(defaultValue: DefaultValue.Custom, defaultValueProvider: provider);
 
-        Assert.AreEqual(provider, Mock.Get(mock).DefaultValueProvider);
+        Assert.AreEqual(provider, mock.DefaultValueProvider);
     }
 
     [TestMethod]
@@ -24,7 +24,7 @@ public void It_uses_default_value_provider_for_a_self_mock_with_interface()
 
         var mock = mocker.WithSelfMock<IService2, Service2>();
 
-        Assert.AreEqual(provider, Mock.Get(mock).DefaultValueProvider);
+        Assert.AreEqual(provider, mock.DefaultValueProvider);
     }
 
     [TestMethod]
@@ -98,5 +98,14 @@ public void It_uses_default_value_provider_for_a_self_mock_with_class_type()
 
         Assert.AreEqual(provider, Mock.Get(mock).DefaultValueProvider);
     }
+
+    [TestMethod]
+    public void WithSelfMock_returns_mock_of_service()
+    {
+        var mocker = new AutoMocker();
+        var mockService = mocker.WithSelfMock<IService2, Service2>();
+
+        Assert.IsInstanceOfType<Mock<IService2>>(mockService);
+    }
 }
 
diff --git a/Moq.AutoMock/AutoMocker.cs b/Moq.AutoMock/AutoMocker.cs
index 67eb0b9d..93159cd8 100644
--- a/Moq.AutoMock/AutoMocker.cs
+++ b/Moq.AutoMock/AutoMocker.cs
@@ -304,8 +304,8 @@ public T CreateSelfMock<T>(
     /// <param name="defaultValue">Sets the DefaultValue property on the created Mock.</param>
     /// <param name="defaultValueProvider">The instance that will be used to produce default return values for unexpected invocations.</param>
     /// <param name="callBase">Sets the CallBase property on the created Mock.</param>
-    /// <returns>An instance with virtual and abstract members mocked</returns> 
-    public TImplementation WithSelfMock<TService, TImplementation>(
+    /// <returns>A mock of the service</returns> 
+    public Mock<TService> WithSelfMock<TService, TImplementation>(
         bool enablePrivate = false,
         MockBehavior? mockBehavior = null,
         DefaultValue? defaultValue = null,
@@ -325,7 +325,7 @@ public TImplementation WithSelfMock<TService, TImplementation>(
             typeMap[typeof(TImplementation)] = new MockInstance(selfMock);
             typeMap[typeof(TService)] = new MockInstance(selfMock.As<TService>());
         });
-        return selfMock.Object;
+        return selfMock.As<TService>();
     }
 
     /// <summary>