Merge pull request #2 from codycollier/no-version-any-avail

When no version is specified, return healthy if any version is available
codycollier · Jul 8, 2022 · be5a7c1 · be5a7c1
2 parents a704a54 + 14fed88
commit be5a7c1
Show file tree

Hide file tree

Showing 2 changed files with 93 additions and 7 deletions.
diff --git a/main.go b/main.go
@@ -61,12 +61,23 @@ func checkServableResponse(response *tfproto.GetModelStatusResponse, modelVersio
 		return 11
 	}
 
-	// Get the state for the noted version. If no version, take the first.
+	// Get the state for the noted version. If no version, take any AVAILABLE.
 	var status tfproto.ModelVersionStatus_State
 	statusFound := false
 	if modelVersion == 0 {
-		status = response.ModelVersionStatus[0].State
-		statusFound = true
+		for _, res := range response.ModelVersionStatus {
+			if res.State == tfproto.ModelVersionStatus_AVAILABLE {
+				status = res.State
+				statusFound = true
+				break
+			}
+		}
+		// when no version is specified, and no model with state available is
+		// found, arbitrarily fallback to first (latest?) item in array
+		if !statusFound {
+			status = response.ModelVersionStatus[0].State
+			statusFound = true
+		}
 	} else {
 		for _, res := range response.ModelVersionStatus {
 			if modelVersion == res.Version {

diff --git a/main_test.go b/main_test.go
@@ -121,23 +121,100 @@ func TestResponseStateAvailable(t *testing.T) {
 	assert.Equal(t, 0, retval, "Expecting response code for state Available")
 }
 
-func TestResponseStateAvailableOnSpecificVersion(t *testing.T) {
+func TestResponseStateAvailableOnNoVersion(t *testing.T) {
 
+	// Ensure success when arbitrary version is available (ex: after a rollback)
 	request := &tfproto.GetModelStatusResponse{
 		ModelVersionStatus: []*tfproto.ModelVersionStatus{
 			{
 				Version: 101,
 				State:   tfproto.ModelVersionStatus_END,
 			},
+			{
+				Version: 98,
+				State:   tfproto.ModelVersionStatus_END,
+			},
 			{
 				Version: 301,
 				State:   tfproto.ModelVersionStatus_AVAILABLE,
 			},
+			{
+				Version: 303,
+				State:   tfproto.ModelVersionStatus_END,
+			},
 		},
 	}
 	retval := checkServableResponse(request, 0)
+	assert.Equal(t, 0, retval)
+
+	// Ensure order is not relevant
+	request = &tfproto.GetModelStatusResponse{
+		ModelVersionStatus: []*tfproto.ModelVersionStatus{
+			{
+				Version: 101,
+				State:   tfproto.ModelVersionStatus_AVAILABLE,
+			},
+			{
+				Version: 301,
+				State:   tfproto.ModelVersionStatus_END,
+			},
+		},
+	}
+	retval = checkServableResponse(request, 0)
+	assert.Equal(t, 0, retval)
+
+	request = &tfproto.GetModelStatusResponse{
+		ModelVersionStatus: []*tfproto.ModelVersionStatus{
+			{
+				Version: 301,
+				State:   tfproto.ModelVersionStatus_END,
+			},
+			{
+				Version: 101,
+				State:   tfproto.ModelVersionStatus_AVAILABLE,
+			},
+		},
+	}
+	retval = checkServableResponse(request, 0)
+	assert.Equal(t, 0, retval)
+
+	// Ensure still fails if there is no version which is available
+	request = &tfproto.GetModelStatusResponse{
+		ModelVersionStatus: []*tfproto.ModelVersionStatus{
+			{
+				Version: 101,
+				State:   tfproto.ModelVersionStatus_END,
+			},
+			{
+				Version: 301,
+				State:   tfproto.ModelVersionStatus_END,
+			},
+			{
+				Version: 303,
+				State:   tfproto.ModelVersionStatus_END,
+			},
+		},
+	}
+	retval = checkServableResponse(request, 0)
 	assert.Equal(t, 34, retval)
-	retval = checkServableResponse(request, 101)
+
+}
+
+func TestResponseStateAvailableOnSpecificVersion(t *testing.T) {
+
+	request := &tfproto.GetModelStatusResponse{
+		ModelVersionStatus: []*tfproto.ModelVersionStatus{
+			{
+				Version: 101,
+				State:   tfproto.ModelVersionStatus_END,
+			},
+			{
+				Version: 301,
+				State:   tfproto.ModelVersionStatus_AVAILABLE,
+			},
+		},
+	}
+	retval := checkServableResponse(request, 101)
 	assert.Equal(t, 34, retval)
 	retval = checkServableResponse(request, 301)
 	assert.Equal(t, 0, retval)
@@ -154,8 +231,6 @@ func TestResponseStateAvailableOnSpecificVersion(t *testing.T) {
 			},
 		},
 	}
-	retval = checkServableResponse(request, 0)
-	assert.Equal(t, 0, retval)
 	retval = checkServableResponse(request, 101)
 	assert.Equal(t, 34, retval)
 	retval = checkServableResponse(request, 301)