Don't use doReturn

airbnb · Feb 28, 2024 · 38206fd · 38206fd
1 parent 2048bb2
commit 38206fd
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/online/src/test/scala/ai/chronon/online/FetcherBaseTest.scala b/online/src/test/scala/ai/chronon/online/FetcherBaseTest.scala
@@ -36,7 +36,7 @@ import scala.concurrent.{Await, ExecutionContext, Future}
 import scala.util.{Failure, Success}
 import scala.util.Try
 
-class FetcherBaseTest extends MockitoSugar with Matchers with MockitoHelper {
+class FetcherBaseTest extends MockitoSugar with Matchers {
   val GroupBy = "relevance.short_term_user_features"
   val Column = "pdp_view_count_14d"
   val GuestKey = "guest"
@@ -153,7 +153,7 @@ class FetcherBaseTest extends MockitoSugar with Matchers with MockitoHelper {
   def test_getServingInfo_ShouldCallUpdateServingInfoIfBatchResponseIsFromKvStore(): Unit = {
     val oldServingInfo = mock[GroupByServingInfoParsed]
     val updatedServingInfo = mock[GroupByServingInfoParsed]
-    doReturn(updatedServingInfo).when(fetcherBase).updateServingInfo(any(), any())
+    when(fetcherBase.updateServingInfo(any(), any())).thenReturn(updatedServingInfo)
 
     val batchTimedValuesSuccess = Success(Seq(TimedValue(Array(1.toByte), 2000L)))
     val kvStoreBatchResponses = BatchResponses(batchTimedValuesSuccess)
@@ -170,16 +170,16 @@ class FetcherBaseTest extends MockitoSugar with Matchers with MockitoHelper {
   @Test
   def test_getServingInfo_ShouldRefreshServingInfoIfBatchResponseIsCached(): Unit = {
     val ttlCache = mock[TTLCache[String, Try[GroupByServingInfoParsed]]]
-    doReturn(ttlCache).when(fetcherBase).getGroupByServingInfo
+    when(fetcherBase.getGroupByServingInfo).thenReturn(ttlCache)
 
     val oldServingInfo = mock[GroupByServingInfoParsed]
-    doReturn(Success(oldServingInfo)).when(ttlCache).refresh(any[String])
+    when(ttlCache.refresh(any[String])).thenReturn(Success(oldServingInfo))
 
     val metaDataMock = mock[MetaData]
     val groupByOpsMock = mock[GroupByOps]
-    doReturn("test").when(metaDataMock).name
-    doReturn(metaDataMock).when(groupByOpsMock).metaData
-    doReturn(groupByOpsMock).when(oldServingInfo).groupByOps
+    when(metaDataMock.name).thenReturn("test")
+    when(groupByOpsMock.metaData).thenReturn(metaDataMock)
+    when(oldServingInfo.groupByOps).thenReturn(groupByOpsMock)
 
     val cachedBatchResponses = BatchResponses(mock[FinalBatchIr])
     val result = fetcherBase.getServingInfo(oldServingInfo, cachedBatchResponses)