Fix channel directions when using experimental modelgen

#79
petergtz · Mar 5, 2019 · cab79a9 · cab79a9
1 parent f35d310
commit cab79a9
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 3 deletions.
diff --git a/dsl_test.go b/dsl_test.go
@@ -882,6 +882,13 @@ var _ = Describe("MockDisplay", func() {
 		})
 	})
 
+	Context("using send-/receive-only channels", func() {
+		It("generates the mock method with correct channel directions", func() {
+			var stringReadChan <-chan string
+			var errorWriteChan chan<- error
+			display.ChanParams(stringReadChan, errorWriteChan)
+		})
+	})
 })
 
 func flattenStringSliceOfSlices(sliceOfSlices [][]string) (result []string) {

diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
@@ -478,7 +478,14 @@ func spaceSeparatedNameFor(t model.Type, packageMap map[string]string) string {
 	case *model.MapType:
 		return "map of " + spaceSeparatedNameFor(typedType.Key, packageMap) + " to " + spaceSeparatedNameFor(typedType.Value, packageMap)
 	case *model.ChanType:
-		return "chan of " + spaceSeparatedNameFor(typedType.Type, packageMap)
+		switch typedType.Dir {
+		case model.RecvDir:
+			return "recv chan of " + spaceSeparatedNameFor(typedType.Type, packageMap)
+		case model.SendDir:
+			return "send chan of " + spaceSeparatedNameFor(typedType.Type, packageMap)
+		default:
+			return "chan of " + spaceSeparatedNameFor(typedType.Type, packageMap)
+		}
 	// TODO:
 	// case *model.FuncType:
 	default:

diff --git a/mockgen/mockgen_test.go b/mockgen/mockgen_test.go
@@ -16,7 +16,7 @@ var _ = Describe("Mockgen", func() {
 			_, matcherSourceCodes := mockgen.GenerateOutput(ast, "irrelevant", "MockDisplay", "test_package", "")
 
 			Expect(matcherSourceCodes).To(SatisfyAll(
-				HaveLen(7),
+				HaveLen(9),
 				HaveKeyWithValue("http_request", SatisfyAll(
 					ContainSubstring("http \"net/http\""),
 					ContainSubstring("func AnyHttpRequest() http.Request"),
@@ -42,6 +42,12 @@ var _ = Describe("Mockgen", func() {
 					ContainSubstring("time \"time\""),
 					ContainSubstring("func AnyTimeTime() time.Time"),
 				)),
+				HaveKeyWithValue("recv_chan_of_string", SatisfyAll(
+					ContainSubstring("func AnyRecvChanOfString() <-chan string"),
+				)),
+				HaveKeyWithValue("send_chan_of_error", SatisfyAll(
+					ContainSubstring("func AnySendChanOfError() chan<- error"),
+				)),
 			))
 		})
 	})

diff --git a/modelgen/loader/loader.go b/modelgen/loader/loader.go
@@ -129,8 +129,17 @@ func (g *modelGenerator) modelTypeFrom(typesType types.Type) model.Type {
 			Value: g.modelTypeFrom(typedTyp.Elem()),
 		}
 	case *types.Chan:
+		var dir model.ChanDir
+		switch typedTyp.Dir() {
+		case types.SendOnly:
+			dir = model.SendDir
+		case types.RecvOnly:
+			dir = model.RecvDir
+		default:
+			dir = 0
+		}
 		return &model.ChanType{
-			Dir:  model.ChanDir(typedTyp.Dir()),
+			Dir:  dir,
 			Type: g.modelTypeFrom(typedTyp.Elem()),
 		}
 	case *types.Named:

diff --git a/test_interface/display.go b/test_interface/display.go
@@ -42,4 +42,5 @@ type Display interface {
 	CamelCaseTypeParam(camelCaseParam io.ReadCloser)
 	MapOfStringToInterfaceParam(m map[string]interface{})
 	UseTime(t time.Time)
+	ChanParams(<-chan string, chan<- error)
 }