Add support for setting the repeat count (#30)

SafetyCulture · Jul 19, 2019 · bc7168f · bc7168f
1 parent 1acca46
commit bc7168f
Show file tree

Hide file tree

Showing 6 changed files with 463 additions and 479 deletions.
diff --git a/protobuf/protoc-gen-gogrpcmock/example/complete.mock.go b/protobuf/protoc-gen-gogrpcmock/example/complete.mock.go
diff --git a/protobuf/protoc-gen-gogrpcmock/example/complete.pb.go b/protobuf/protoc-gen-gogrpcmock/example/complete.pb.go
diff --git a/protobuf/protoc-gen-gogrpcmock/example/complete.proto b/protobuf/protoc-gen-gogrpcmock/example/complete.proto
@@ -27,7 +27,7 @@ message Complete {
   repeated string word = 7;
   string url = 8;
   int32 single_number = 9;
-  repeated int64 repeated_number = 10;
+  repeated int64 repeated_number = 10 [ (grpcmock.field) = {repeatn : 3} ];
   int32 lat = 11;
   int32 lng = 12;
   string words = 13 [ (grpcmock.field) = {words : true} ];

diff --git a/protobuf/protoc-gen-gogrpcmock/plugin/plugin.go b/protobuf/protoc-gen-gogrpcmock/plugin/plugin.go
@@ -213,7 +213,7 @@ func (g *grpcmock) generateMockString(fieldName, fieldType string, repeated bool
 		g.P(fieldName, `: `, fieldType, `{`)
 		g.In()
 
-		for i := 0; i < repeatCount; i++ {
+		for i := 0; i < getRepeatCount(field); i++ {
 			g.P(`"`, generateStringValue(fieldName, field), `",`)
 		}
 
@@ -231,7 +231,7 @@ func (g *grpcmock) generateMockInt(fieldName, fieldType string, repeated bool, f
 		g.P(fieldName, `: `, fieldType, `{`)
 		g.In()
 
-		for i := 0; i < repeatCount; i++ {
+		for i := 0; i < getRepeatCount(field); i++ {
 			g.P(generateIntValue(fieldName, field), `,`)
 		}
 
@@ -247,7 +247,7 @@ func (g *grpcmock) generateMockFloat(fieldName, fieldType string, repeated bool,
 		g.P(fieldName, `: `, fieldType, `{`)
 		g.In()
 
-		for i := 0; i < repeatCount; i++ {
+		for i := 0; i < getRepeatCount(field); i++ {
 			g.P(generateFloatValue(fieldName, field), `,`)
 		}
 
@@ -286,7 +286,7 @@ func (g *grpcmock) generateMockInnerMessage(fieldName, fieldType string, repeate
 	length := 1
 
 	if repeated {
-		length = repeatCount
+		length = getRepeatCount(field)
 		g.P(fieldName, `: `, fieldType, `{`)
 		g.In()
 	} else {
@@ -298,7 +298,9 @@ func (g *grpcmock) generateMockInnerMessage(fieldName, fieldType string, repeate
 	for i := 0; i < length; i++ {
 		switch fieldType {
 		case "time.Time":
-			g.P(`time.Now(),`)
+			g.P(`time.Now().Add(` + strconv.Itoa(r.Intn(1000)) + `),`)
+		case "time.Duration":
+			g.P(`time.Duration(` + strconv.Itoa(r.Intn(1000)) + `),`)
 		case "*time.Time", "*time.Duration":
 			// g.P(`&time.Time{},`)
 			g.P(`nil,`)
@@ -470,7 +472,7 @@ func generateIntValue(fieldName string, field *descriptor.FieldDescriptorProto)
 		}
 	}
 
-	return strconv.Itoa(int(r.Intn(n)))
+	return strconv.Itoa(r.Intn(n))
 }
 
 func generateFloatValue(fieldName string, field *descriptor.FieldDescriptorProto) string {
@@ -507,3 +509,10 @@ func boolFromPtr(b *bool) bool {
 	}
 	return *b
 }
+
+func getRepeatCount(field *descriptor.FieldDescriptorProto) int {
+	if mocks := getFieldMocksIfAny(field); mocks != nil && mocks.Repeatn != nil {
+		return int(mocks.GetRepeatn())
+	}
+	return repeatCount
+}
diff --git a/protobuf/s12proto/grpcmock.pb.go b/protobuf/s12proto/grpcmock.pb.go
diff --git a/protobuf/s12proto/grpcmock.proto b/protobuf/s12proto/grpcmock.proto
@@ -33,4 +33,5 @@ message FieldMock {
   optional bool hexcolor = 19;
   repeated string prefix = 20;
   optional int32 floatn = 21;
+  optional int32 repeatn = 22;
 }