add unit test + fix queryModels

EXAMPLES.md

@@ -1176,7 +1176,7 @@ peer chaincode query -n mycc -c '{"Args":["queryModels"]}' -C myc
 ##### Command output:
 ```json
 {
- "bookmark": "",
+ "bookmark": "{\"traintuple\":\"\",\"composite_traintuple\":\"\",\"aggregatetuple\":\"\"}",
  "results": [
   {
    "traintuple": {

chaincode/input.go

@@ -123,7 +123,7 @@ type inputKey struct {
 }
 
 type inputBookmark struct {
-	Bookmark string `validate:"required" json:"bookmark"`
+	Bookmark string `json:"bookmark"`
 }
 
 type inputLogSuccessTrain struct {

chaincode/main_test.go

@@ -342,9 +342,17 @@ func TestQueryEmptyResponse(t *testing.T) {
 			args := [][]byte{[]byte(contractName)}
 			resp := mockStub.MockInvoke(args)
 
+			expectedBookmark := ""
+			if contractName == "queryModels" {
+				// special case for queryModels: we return a map instead of a string
+				expectedBookmarkBytes, _ := json.Marshal(queryModelsBookmarks{})
+				expectedBookmark = string(expectedBookmarkBytes)
+			}
+
 			expectedResult := map[string]interface{}{
 				"results":  make([]string, 0),
-				"bookmark": ""}
+				"bookmark": expectedBookmark}
+
 			expectedPayload, _ := json.Marshal(expectedResult)
 			assert.Equal(t, expectedPayload, resp.Payload, "payload is not an empty list")
 		})

chaincode/mockstub_test.go

@@ -283,7 +283,7 @@ func (stub *MockStub) GetStateByRange(startKey, endKey string) (shim.StateQueryI
 	if err := validateSimpleKeys(startKey, endKey); err != nil {
 		return nil, err
 	}
-	return NewMockStateRangeQueryIterator(stub, startKey, endKey), nil
+	return NewMockStateRangeQueryIterator(stub, startKey, endKey, false, OutputPageSize), nil
 }
 
 // GetQueryResult function can be invoked by a chaincode to perform a
@@ -315,7 +315,7 @@ func (stub *MockStub) GetStateByPartialCompositeKey(objectType string, attribute
 	if err != nil {
 		return nil, err
 	}
-	return NewMockStateRangeQueryIterator(stub, partialCompositeKey, partialCompositeKey+string(maxUnicodeRuneValue)), nil
+	return NewMockStateRangeQueryIterator(stub, partialCompositeKey, partialCompositeKey+string(maxUnicodeRuneValue), false, OutputPageSize), nil
 }
 
 // CreateCompositeKey combines the list of attributes
@@ -341,7 +341,12 @@ func (stub *MockStub) GetStateByPartialCompositeKeyWithPagination(objectType str
 	if err != nil {
 		return nil, nil, err
 	}
-	return NewMockStateRangeQueryIterator(stub, partialCompositeKey, partialCompositeKey+string(maxUnicodeRuneValue)), nil, nil
+	startKey := partialCompositeKey
+	if bookmark != "" {
+		startKey = bookmark
+	}
+	iterator := NewMockStateRangeQueryIterator(stub, startKey, partialCompositeKey+string(maxUnicodeRuneValue), true, pageSize)
+	return iterator, iterator.Metadata, nil
 }
 
 func (stub *MockStub) GetQueryResultWithPagination(query string, pageSize int32,
@@ -490,7 +495,7 @@ func NewMockStub(name string, cc shim.Chaincode) *MockStub {
 	s.EndorsementPolicies = make(map[string]map[string][]byte)
 	s.Invokables = make(map[string]*MockStub)
 	s.Keys = list.New()
-	s.ChaincodeEventsChannel = make(chan *pb.ChaincodeEvent, 100) //define large capacity for non-blocking setEvent calls.
+	s.ChaincodeEventsChannel = make(chan *pb.ChaincodeEvent, OutputPageSize+1) //define large capacity for non-blocking setEvent calls.
 	s.Decorations = make(map[string][]byte)
 	s.TxTimestamp = &timestamp.Timestamp{}
 
@@ -509,11 +514,14 @@ func NewMockStubWithRegisterNode(name string, cc shim.Chaincode) *MockStub {
 *****************************/
 
 type MockStateRangeQueryIterator struct {
-	Closed   bool
-	Stub     *MockStub
-	StartKey string
-	EndKey   string
-	Current  *list.Element
+	Closed      bool
+	Stub        *MockStub
+	StartKey    string
+	EndKey      string
+	Current     *list.Element
+	IsPaginated bool
+	PageSize    int32
+	Metadata    *pb.QueryResponseMetadata
 }
 
 // HasNext returns true if the range query iterator contains additional keys
@@ -534,12 +542,17 @@ func (iter *MockStateRangeQueryIterator) HasNext() bool {
 		if iter.StartKey == "" && iter.EndKey == "" {
 			return true
 		}
+		// iterator has already yielded enough results
+		if iter.IsPaginated && iter.Metadata.FetchedRecordsCount == iter.PageSize {
+			return false
+		}
 		comp1 := strings.Compare(current.Value.(string), iter.StartKey)
 		comp2 := strings.Compare(current.Value.(string), iter.EndKey)
 		if comp1 >= 0 {
 			if comp2 < 0 {
 				return true
 			}
+			iter.Metadata.Bookmark = ""
 			return false
 		}
 		current = current.Next()
@@ -561,6 +574,10 @@ func (iter *MockStateRangeQueryIterator) Next() (*queryresult.KV, error) {
 		return nil, err
 	}
 
+	if iter.Current != nil && iter.IsPaginated && iter.Metadata.FetchedRecordsCount >= iter.PageSize {
+		return nil, errors.New("Paginated MockStateRangeQueryIterator.Next() went past end of range")
+	}
+
 	for iter.Current != nil {
 		comp1 := strings.Compare(iter.Current.Value.(string), iter.StartKey)
 		comp2 := strings.Compare(iter.Current.Value.(string), iter.EndKey)
@@ -570,6 +587,12 @@ func (iter *MockStateRangeQueryIterator) Next() (*queryresult.KV, error) {
 			key := iter.Current.Value.(string)
 			value, err := iter.Stub.GetState(key)
 			iter.Current = iter.Current.Next()
+			if iter.Current != nil {
+				iter.Metadata.Bookmark = iter.Current.Value.(string)
+			} else {
+				iter.Metadata.Bookmark = ""
+			}
+			iter.Metadata.FetchedRecordsCount++
 			return &queryresult.KV{Key: key, Value: value}, err
 		}
 		iter.Current = iter.Current.Next()
@@ -593,13 +616,16 @@ func (iter *MockStateRangeQueryIterator) Close() error {
 func (iter *MockStateRangeQueryIterator) Print() {
 }
 
-func NewMockStateRangeQueryIterator(stub *MockStub, startKey string, endKey string) *MockStateRangeQueryIterator {
+func NewMockStateRangeQueryIterator(stub *MockStub, startKey string, endKey string, isPaginated bool, pageSize int32) *MockStateRangeQueryIterator {
 	iter := new(MockStateRangeQueryIterator)
 	iter.Closed = false
 	iter.Stub = stub
 	iter.StartKey = startKey
 	iter.EndKey = endKey
 	iter.Current = stub.Keys.Front()
+	iter.IsPaginated = isPaginated
+	iter.PageSize = pageSize
+	iter.Metadata = &pb.QueryResponseMetadata{}
 
 	iter.Print()
 

chaincode/output.go

@@ -20,7 +20,7 @@ import (
 )
 
 // OutputPageSize is a used to avoid issues listing assets
-const OutputPageSize = 5 // 500
+const OutputPageSize = 500
 
 // Struct use as output representation of ledger data
 

chaincode/traintuple_test.go

@@ -56,6 +56,53 @@ func TestTraintupleWithNoTestDataset(t *testing.T) {
 	assert.EqualValues(t, 200, resp.Status, "It should find the traintuple without error ", resp.Message)
 }
 
+func TestPagination(t *testing.T) {
+	scc := new(SubstraChaincode)
+	mockStub := NewMockStubWithRegisterNode("substra", scc)
+	registerItem(t, *mockStub, "trainDataset")
+
+	key := strings.Replace(objectiveKey, "1", "2", 1)
+	inpObjective := inputObjective{Key: key}
+	inpObjective.createDefault()
+	inpObjective.TestDataset = inputDataset{}
+	resp := mockStub.MockInvoke(methodAndAssetToByte("registerObjective", inpObjective))
+	assert.EqualValues(t, 200, resp.Status, "when adding objective without dataset it should work: ", resp.Message)
+
+	inpAlgo := inputAlgo{}
+	args := inpAlgo.createDefault()
+	resp = mockStub.MockInvoke(args)
+	assert.EqualValues(t, 200, resp.Status, "when adding algo it should work: ", resp.Message)
+
+	// Add N + 1 testtuples
+	for i := 0; i < OutputPageSize+1; i++ {
+		uuid, err := GetNewUUID()
+		assert.NoError(t, err)
+		inpTraintuple := inputTraintuple{Key: uuid}
+		args = inpTraintuple.createDefault()
+		resp = mockStub.MockInvoke(args)
+		assert.EqualValues(t, 200, resp.Status, "when adding traintuple without test dataset it should work: ", resp.Message)
+	}
+
+	var queryTraintuples TraintupleResponse
+
+	// 1st query (no bookmark) should return OutputPageSize results
+	args = [][]byte{[]byte("queryTraintuples")}
+	resp = mockStub.MockInvoke(args)
+	assert.EqualValues(t, 200, resp.Status, "It should find the traintuple without error ", resp.Message)
+	err := json.Unmarshal(resp.Payload, &queryTraintuples)
+	assert.NoError(t, err, "traintuples should unmarshal without problem")
+	assert.Equal(t, OutputPageSize, len(queryTraintuples.Results))
+
+	// 2nd query (with bookmark) should return 1 result
+	inp := inputBookmark{Bookmark: queryTraintuples.Bookmark}
+	args = append([][]byte{[]byte("queryTraintuples")}, assetToJSON(inp))
+	resp = mockStub.MockInvoke(args)
+	assert.EqualValues(t, 200, resp.Status, "It should find the traintuple without error ", resp.Message)
+	err = json.Unmarshal(resp.Payload, &queryTraintuples)
+	assert.NoError(t, err, "traintuples should unmarshal without problem")
+	assert.Equal(t, 1, len(queryTraintuples.Results))
+}
+
 func TestTraintupleWithSingleDatasample(t *testing.T) {
 	scc := new(SubstraChaincode)
 	mockStub := NewMockStubWithRegisterNode("substra", scc)

chaincode/tuple.go

@@ -95,26 +95,43 @@ func queryModelDetails(db *LedgerDB, args []string) (outModelDetails outputModel
 	return
 }
 
+type queryModelsBookmarks struct {
+	Traintuple          string `json:"traintuple"`
+	CompositeTraintuple string `json:"composite_traintuple"`
+	Aggregatetuple      string `json:"aggregatetuple"`
+}
+
+type inputQueryModelsBookmarks struct {
+	Bookmark string `json:"bookmark"`
+}
+
 // queryModels returns all traintuples and associated testuples
 func queryModels(db *LedgerDB, args []string) (outModels []outputModel, bookmark string, err error) {
 	outModels = []outputModel{}
-	bookmarks := map[string]string{}
+	bookmarks := queryModelsBookmarks{}
 
 	if len(args) > 1 {
 		err = errors.BadRequest("incorrect number of arguments, expecting at most one argument")
 		return
 	}
 
 	if len(args) == 1 && args[0] != "" {
-		err = json.Unmarshal([]byte(args[0]), &bookmarks)
+		inp := inputQueryModelsBookmarks{}
+		err = json.Unmarshal([]byte(args[0]), &inp)
 		if err != nil {
 			return
 		}
+		if inp.Bookmark != "" {
+			err = json.Unmarshal([]byte(inp.Bookmark), &bookmarks)
+			if err != nil {
+				return
+			}
+		}
 	}
 
 	// populate from regular traintuples
-	traintupleKeys, traintupleBookmark, err := db.GetIndexKeysWithPagination("traintuple~algo~key", []string{"traintuple"}, OutputPageSize/3, bookmarks["traintuple"])
-	bookmarks["traintuple"] = traintupleBookmark
+	traintupleKeys, traintupleBookmark, err := db.GetIndexKeysWithPagination("traintuple~algo~key", []string{"traintuple"}, OutputPageSize/3, bookmarks.Traintuple)
+	bookmarks.Traintuple = traintupleBookmark
 
 	if err != nil {
 		return
@@ -133,8 +150,8 @@ func queryModels(db *LedgerDB, args []string) (outModels []outputModel, bookmark
 	}
 
 	// populate from composite traintuples
-	compositeTraintupleKeys, compositeTraintupleBookmark, err := db.GetIndexKeysWithPagination("compositeTraintuple~algo~key", []string{"compositeTraintuple"}, OutputPageSize/3, bookmarks["compositeTraintuple"])
-	bookmarks["compositeTraintuple"] = compositeTraintupleBookmark
+	compositeTraintupleKeys, compositeTraintupleBookmark, err := db.GetIndexKeysWithPagination("compositeTraintuple~algo~key", []string{"compositeTraintuple"}, OutputPageSize/3, bookmarks.CompositeTraintuple)
+	bookmarks.CompositeTraintuple = compositeTraintupleBookmark
 
 	if err != nil {
 		return
@@ -152,8 +169,8 @@ func queryModels(db *LedgerDB, args []string) (outModels []outputModel, bookmark
 	}
 
 	// populate from composite traintuples
-	aggregatetupleKeys, aggregatetupleBookmark, err := db.GetIndexKeysWithPagination("aggregatetuple~algo~key", []string{"aggregatetuple"}, OutputPageSize/3, bookmarks["aggregateTuple"])
-	bookmarks["aggregateTuple"] = aggregatetupleBookmark
+	aggregatetupleKeys, aggregatetupleBookmark, err := db.GetIndexKeysWithPagination("aggregatetuple~algo~key", []string{"aggregatetuple"}, OutputPageSize/3, bookmarks.Aggregatetuple)
+	bookmarks.Aggregatetuple = aggregatetupleBookmark
 
 	if err != nil {
 		return
@@ -170,6 +187,8 @@ func queryModels(db *LedgerDB, args []string) (outModels []outputModel, bookmark
 		outModels = append(outModels, outputModel)
 	}
 
+	bookmarkBytes, err := json.Marshal(bookmarks)
+	bookmark = string(bookmarkBytes)
 	return
 }
 

chaincode/tuple_test.go

@@ -222,3 +222,57 @@ func TestQueryHeadModelPermissions(t *testing.T) {
 	assert.Len(t, outPerm.Process.AuthorizedIDs, 1)
 	assert.Contains(t, outPerm.Process.AuthorizedIDs, worker)
 }
+
+type ModelsResponse struct {
+	Results  []outputModel `json:"results"`
+	Bookmark string        `json:"bookmark"`
+}
+
+func TestQueryModelsPagination(t *testing.T) {
+	scc := new(SubstraChaincode)
+	mockStub := NewMockStubWithRegisterNode("substra", scc)
+
+	resp, _ := registerItem(t, *mockStub, "algo")
+
+	// Add N + 1 traintuples
+	for i := 0; i < OutputPageSize/3+1; i++ {
+		uuid, _ := GetNewUUID()
+		inpTraintuple := inputTraintuple{Key: uuid}
+		args := inpTraintuple.createDefault()
+		resp = mockStub.MockInvoke(args)
+		assert.EqualValues(t, 200, resp.Status, resp.Message)
+
+		args = [][]byte{[]byte("logStartTrain"), keyToJSON(inpTraintuple.Key)}
+		resp = mockStub.MockInvoke(args)
+
+		success := inputLogSuccessTrain{}
+		success.Key = traintupleKey
+		success.OutModel.Key, _ = GetNewUUID()
+		args = append([][]byte{[]byte("logSuccessTrain")}, assetToJSON(success))
+		resp = mockStub.MockInvoke(args)
+	}
+
+	var models ModelsResponse
+
+	// 1st query (no bookmark) should return OutputPageSize/3 results
+	args := [][]byte{[]byte("queryModels")}
+	resp = mockStub.MockInvoke(args)
+	assert.EqualValues(t, 200, resp.Status, "It should find the models without error ", resp.Message)
+	err := json.Unmarshal(resp.Payload, &models)
+	assert.NoError(t, err, "models should unmarshal without problem")
+	assert.Equal(t, OutputPageSize/3, len(models.Results))
+	firstResult := models.Results[0].Traintuple.Key
+
+	// 2nd query (with bookmark) should return 1 result
+	inp := inputBookmark{Bookmark: models.Bookmark}
+	args = append([][]byte{[]byte("queryModels")}, assetToJSON(inp))
+	resp = mockStub.MockInvoke(args)
+	assert.EqualValues(t, 200, resp.Status, "It should find the models without error ", resp.Message)
+	err = json.Unmarshal(resp.Payload, &models)
+	assert.NoError(t, err, "models should unmarshal without problem")
+	assert.Equal(t, 1, len(models.Results))
+
+	// 2nd query should return different results from 1st query
+	newFirstResult := models.Results[0].Traintuple.Key
+	assert.NotEqual(t, newFirstResult, firstResult, "query results should be different")
+}