From 73672b2d516875312561b4b91e0cec8c210089f4 Mon Sep 17 00:00:00 2001
From: tyiuhc <tim.yiu@amplitude.com>
Date: Tue, 30 Jul 2024 17:02:57 -0700
Subject: [PATCH] Cohort not modified should not throw exception

---
 pkg/experiment/local/cohort_download_api.go      | 4 ++--
 pkg/experiment/local/cohort_download_api_test.go | 7 +++----
 pkg/experiment/local/cohort_loader.go            | 4 +++-
 pkg/experiment/local/exception.go                | 8 --------
 4 files changed, 8 insertions(+), 15 deletions(-)

diff --git a/pkg/experiment/local/cohort_download_api.go b/pkg/experiment/local/cohort_download_api.go
index 45e5320..6d09fcd 100644
--- a/pkg/experiment/local/cohort_download_api.go
+++ b/pkg/experiment/local/cohort_download_api.go
@@ -48,7 +48,7 @@ func (api *directCohortDownloadApi) getCohort(cohortID string, cohort *Cohort) (
 			errors++
 			if errors >= 3 || func(err error) bool {
 				switch err.(type) {
-				case *CohortNotModifiedException, *CohortTooLargeException:
+				case *CohortTooLargeException:
 					return true
 				default:
 					return false
@@ -85,7 +85,7 @@ func (api *directCohortDownloadApi) getCohort(cohortID string, cohort *Cohort) (
 				}(),
 			}, nil
 		} else if response.StatusCode == http.StatusNoContent {
-			return nil, &CohortNotModifiedException{Message: "Cohort not modified"}
+			return nil, nil
 		} else if response.StatusCode == http.StatusRequestEntityTooLarge {
 			return nil, &CohortTooLargeException{Message: "Cohort exceeds max cohort size"}
 		} else {
diff --git a/pkg/experiment/local/cohort_download_api_test.go b/pkg/experiment/local/cohort_download_api_test.go
index 8029283..fbd6422 100644
--- a/pkg/experiment/local/cohort_download_api_test.go
+++ b/pkg/experiment/local/cohort_download_api_test.go
@@ -214,9 +214,8 @@ func TestCohortDownloadApi(t *testing.T) {
 			httpmock.NewStringResponder(204, ""),
 		)
 
-		_, err := api.getCohort("1234", cohort)
-		assert.Error(t, err)
-		_, isCohortNotModifiedException := err.(*CohortNotModifiedException)
-		assert.True(t, isCohortNotModifiedException)
+		result, err := api.getCohort("1234", cohort)
+		assert.Nil(t, result)
+		assert.NoError(t, err)
 	})
 }
diff --git a/pkg/experiment/local/cohort_loader.go b/pkg/experiment/local/cohort_loader.go
index 7023306..94d30c0 100644
--- a/pkg/experiment/local/cohort_loader.go
+++ b/pkg/experiment/local/cohort_loader.go
@@ -69,7 +69,9 @@ func (task *CohortLoaderTask) run() {
 	if err != nil {
 		task.err = err
 	} else {
-		task.loader.cohortStorage.putCohort(cohort)
+		if cohort != nil {
+			task.loader.cohortStorage.putCohort(cohort)
+		}
 	}
 
 	task.loader.removeJob(task.cohortId)
diff --git a/pkg/experiment/local/exception.go b/pkg/experiment/local/exception.go
index c151916..498d3ab 100644
--- a/pkg/experiment/local/exception.go
+++ b/pkg/experiment/local/exception.go
@@ -16,11 +16,3 @@ type CohortTooLargeException struct {
 func (e *CohortTooLargeException) Error() string {
 	return e.Message
 }
-
-type CohortNotModifiedException struct {
-	Message string
-}
-
-func (e *CohortNotModifiedException) Error() string {
-	return e.Message
-}