Context("when there are more absences of the target key than occurances", func() { It("decides with a value of false", func() { keyCountCall.Return(3) mockBinaryTree.EXPECT().Set(false) myTree.Train(inputLog, targetKey) }) }) }) Describe("when there are still features to decide on", func() { BeforeEach(func() { nextKey = "the next key to decide on" maxEntropyCall = maxEntropyCall.Return(nextKey) leftDecisionTree = mdt.NewMockInterface(mockCtrl) rightDecisionTree = mdt.NewMockInterface(mockCtrl) decisionTreeMocks := []*mdt.MockInterface{rightDecisionTree, leftDecisionTree} // Inject left and right tree in alternating fashion leftRightParity := 0 dt.New = func() dt.Interface { leftRightParity++ return decisionTreeMocks[leftRightParity%2] } inputLog.EXPECT().SplitOnKey(nextKey).Return(leftInputLog, rightInputLog) mockBinaryTree.EXPECT().InsertLeft(leftDecisionTree) mockBinaryTree.EXPECT().InsertRight(rightDecisionTree) keySetCall = mockBinaryTree.EXPECT().Set(nextKey).AnyTimes()
mockCtrl = gomock.NewController(gomocktestreporter.New()) }) AfterEach(func() { dt.New = decisionTreeNew mockCtrl.Finish() }) Describe("Train", func() { var ( predictionForestInputSet il.InputLog mockDecisionTree *mdt.MockInterface ) BeforeEach(func() { mockDecisionTree = mdt.NewMockInterface(mockCtrl) dt.New = func() dt.Interface { return mockDecisionTree } predictionForestInputSet = il.InputLog{ {"1", "3"}, {"0", "1", "3"}, {"1", "3"}, {"0", "1", "3"}, {"1", "2"}, {"0", "1", "2"}, {}, {"0"}, } })