From 960030a48252ae20010f62895db9d1f988592d92 Mon Sep 17 00:00:00 2001
From: Jason Parraga <sovietaced@gmail.com>
Date: Thu, 1 Aug 2024 22:12:30 -0700
Subject: [PATCH] Add unit test for panic scenario

---
 flytestdlib/cache/auto_refresh_test.go | 37 ++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/flytestdlib/cache/auto_refresh_test.go b/flytestdlib/cache/auto_refresh_test.go
index e798300f5d0..59ae4b7494f 100644
--- a/flytestdlib/cache/auto_refresh_test.go
+++ b/flytestdlib/cache/auto_refresh_test.go
@@ -64,6 +64,15 @@ func syncTerminalItem(_ context.Context, batch Batch) ([]ItemSyncResponse, error
 	panic("This should never be called")
 }
 
+type panickingSyncer struct {
+	callCount atomic.Int32
+}
+
+func (p *panickingSyncer) sync(_ context.Context, _ Batch) ([]ItemSyncResponse, error) {
+	p.callCount.Inc()
+	panic("testing")
+}
+
 func TestCacheFour(t *testing.T) {
 	testResyncPeriod := 10 * time.Millisecond
 	rateLimiter := workqueue.DefaultControllerRateLimiter()
@@ -172,6 +181,34 @@ func TestCacheFour(t *testing.T) {
 
 		cancel()
 	})
+
+	t.Run("Test panic on sync and shutdown", func(t *testing.T) {
+		syncer := &panickingSyncer{}
+		cache, err := NewAutoRefreshCache("fake3", syncer.sync, rateLimiter, testResyncPeriod, 10, 2, promutils.NewTestScope())
+		assert.NoError(t, err)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		assert.NoError(t, cache.Start(ctx))
+
+		itemID := "dummy_id"
+		_, err = cache.GetOrCreate(itemID, fakeCacheItem{
+			val: 0,
+		})
+		assert.NoError(t, err)
+
+		// wait for all workers to run
+		assert.Eventually(t, func() bool {
+			return syncer.callCount.Load() == int32(10)
+		}, time.Second, time.Millisecond)
+
+		// wait some more time
+		time.Sleep(500 * time.Millisecond)
+
+		// all workers should have shut down.
+		assert.Equal(t, int32(10), syncer.callCount.Load())
+
+		cancel()
+	})
 }
 
 func TestQueueBuildUp(t *testing.T) {