CosmosDB 1.0 code change

state/azure/cosmosdb/cosmosdb.go

@@ -81,7 +81,6 @@ type CosmosItem struct {
 const (
 	metadataPartitionKey = "partitionKey"
 	defaultTimeout       = 20 * time.Second
-	statusNotFound       = "NotFound"
 )
 
 // Policy that makes all queries cross-partition
@@ -228,8 +227,7 @@ func (c *StateStore) Get(ctx context.Context, req *state.GetRequest) (*state.Get
 	defer cancel()
 	readItem, err := c.client.ReadItem(readCtx, azcosmos.NewPartitionKeyString(partitionKey), req.Key, &options)
 	if err != nil {
-		var responseErr *azcore.ResponseError
-		if errors.As(err, &responseErr) && responseErr.ErrorCode == "NotFound" {
+		if isNotFoundError(err) {
 			return &state.GetResponse{}, nil
 		}
 		return nil, err
@@ -690,9 +688,10 @@ func isNotFoundError(err error) bool {
 	}
 
 	if requestError, ok := err.(*azcore.ResponseError); ok {
-		if requestError.ErrorCode == statusNotFound {
+		if requestError.StatusCode == 404 {
 			return true
 		}
+		// we previously checked the error code, but unfortunately this is not stable between API versions
 	}
 
 	return false

state/azure/cosmosdb/cosmosdb_query.go

@@ -261,7 +261,7 @@
 	opts.QueryParameters = append(opts.QueryParameters, q.query.parameters...)
 
 	if len(q.token) != 0 {
-		opts.ContinuationToken = q.token
+		opts.ContinuationToken = &q.token
 	}
 
 	items := []CosmosItem{}
@@ -285,7 +285,11 @@
 			return nil, "", innerErr
 		}
 
-		token = queryResponse.ContinuationToken
+		if queryResponse.ContinuationToken == nil {
+			token = ""
+		} else {
+			token = *queryResponse.ContinuationToken
+		}
 		for _, item := range queryResponse.Items {
 			tempItem := CosmosItem{}
 			err := json.Unmarshal(item, &tempItem)

tests/certification/state/azure/cosmosdb/cosmosdb_test.go

@@ -133,13 +133,17 @@ func TestAzureCosmosDBStorage(t *testing.T) {
 				"partitionKey": "mypartition",
 			}
 
-			test := func(setMeta, getMeta map[string]string, expectedValue string) {
+			test := func(setMeta, getMeta map[string]string, expectedValue string, expectedErr bool) {
 				// save state, default options: strong, last-write
 				err = client.SaveState(ctx, statestore, stateKey, []byte(stateValue), setMeta)
 				require.NoError(t, err)
 
 				// get state
 				item, err := client.GetState(ctx, statestore, stateKey, getMeta)
+				if expectedErr {
+					require.Error(t, err)
+					return
+				}
 				require.NoError(t, err)
 				assert.Equal(t, expectedValue, string(item.Value))
 
@@ -149,13 +153,13 @@ func TestAzureCosmosDBStorage(t *testing.T) {
 			}
 
 			// Test	with no partition key
-			test(nil, meta1, stateValue)
+			test(nil, meta1, stateValue, false)
 
 			// Test with specific partition key
-			test(meta2, meta2, stateValue)
+			test(meta2, meta2, stateValue, false)
 
 			// Test with incorrect partition key
-			test(meta2, meta1, "")
+			test(meta2, meta1, "", true)
 
 			return nil
 		}