From 2048bb278bce85ac4876a6d3b7de783f20b50f96 Mon Sep 17 00:00:00 2001
From: Caio Camatta <caiocamatta@stripe.com>
Date: Wed, 28 Feb 2024 09:39:41 -0500
Subject: [PATCH] Refactor getServingInfo tests

---
 .../ai/chronon/online/FetcherBaseTest.scala   | 41 ++++++++++---------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/online/src/test/scala/ai/chronon/online/FetcherBaseTest.scala b/online/src/test/scala/ai/chronon/online/FetcherBaseTest.scala
index e033c4888..bc7b60d1b 100644
--- a/online/src/test/scala/ai/chronon/online/FetcherBaseTest.scala
+++ b/online/src/test/scala/ai/chronon/online/FetcherBaseTest.scala
@@ -30,7 +30,6 @@ import org.mockito.stubbing.Answer
 import org.mockito.{Answers, ArgumentCaptor}
 import org.scalatest.matchers.should.Matchers
 import org.scalatestplus.mockito.MockitoSugar
-import org.junit.Assert.assertSame
 
 import scala.concurrent.duration.DurationInt
 import scala.concurrent.{Await, ExecutionContext, Future}
@@ -149,41 +148,45 @@ class FetcherBaseTest extends MockitoSugar with Matchers with MockitoHelper {
     actualRequest.get.keys shouldBe query.keyMapping.get
   }
 
+  // updateServingInfo() is called when the batch response is from the KV store.
   @Test
   def test_getServingInfo_ShouldCallUpdateServingInfoIfBatchResponseIsFromKvStore(): Unit = {
-    val baseFetcher = new FetcherBase(mock[KVStore])
-    val spiedFetcherBase = spy(baseFetcher)
     val oldServingInfo = mock[GroupByServingInfoParsed]
     val updatedServingInfo = mock[GroupByServingInfoParsed]
+    doReturn(updatedServingInfo).when(fetcherBase).updateServingInfo(any(), any())
+
     val batchTimedValuesSuccess = Success(Seq(TimedValue(Array(1.toByte), 2000L)))
     val kvStoreBatchResponses = BatchResponses(batchTimedValuesSuccess)
-    doReturn(updatedServingInfo).when(spiedFetcherBase).updateServingInfo(any(), any())
+
+    val result = fetcherBase.getServingInfo(oldServingInfo, kvStoreBatchResponses)
 
     // updateServingInfo is called
-    val result = spiedFetcherBase.getServingInfo(oldServingInfo, kvStoreBatchResponses)
-    assertSame(result, updatedServingInfo)
-    verify(spiedFetcherBase).updateServingInfo(any(), any())
+    result shouldEqual updatedServingInfo
+    verify(fetcherBase).updateServingInfo(any(), any())
   }
 
+  // If a batch response is cached, the serving info should be refreshed. This is needed to prevent
+  // the serving info from becoming stale if all the requests are cached.
   @Test
   def test_getServingInfo_ShouldRefreshServingInfoIfBatchResponseIsCached(): Unit = {
-    val baseFetcher = new FetcherBase(mock[KVStore])
-    val spiedFetcherBase = spy(baseFetcher)
+    val ttlCache = mock[TTLCache[String, Try[GroupByServingInfoParsed]]]
+    doReturn(ttlCache).when(fetcherBase).getGroupByServingInfo
+
     val oldServingInfo = mock[GroupByServingInfoParsed]
-    val metaData = mock[MetaData]
+    doReturn(Success(oldServingInfo)).when(ttlCache).refresh(any[String])
+
+    val metaDataMock = mock[MetaData]
     val groupByOpsMock = mock[GroupByOps]
+    doReturn("test").when(metaDataMock).name
+    doReturn(metaDataMock).when(groupByOpsMock).metaData
+    doReturn(groupByOpsMock).when(oldServingInfo).groupByOps
+
     val cachedBatchResponses = BatchResponses(mock[FinalBatchIr])
-    val ttlCache = mock[TTLCache[String, Try[GroupByServingInfoParsed]]]
-    doReturn(ttlCache).when(spiedFetcherBase).getGroupByServingInfo
-    doReturn(Success(oldServingInfo)).when(ttlCache).refresh(any[String])
-    metaData.name = "test"
-    groupByOpsMock.metaData = metaData
-    when(oldServingInfo.groupByOps).thenReturn(groupByOpsMock)
+    val result = fetcherBase.getServingInfo(oldServingInfo, cachedBatchResponses)
 
     // FetcherBase.updateServingInfo is not called, but getGroupByServingInfo.refresh() is.
-    val result = spiedFetcherBase.getServingInfo(oldServingInfo, cachedBatchResponses)
-    assertSame(result, oldServingInfo)
+    result shouldEqual oldServingInfo
     verify(ttlCache).refresh(any())
-    verify(spiedFetcherBase, never()).updateServingInfo(any(), any())
+    verify(fetcherBase, never()).updateServingInfo(any(), any())
   }
 }