1 // Copyright 2020 The Chromium Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 #include "components/optimization_guide/optimization_guide_test_util.h"
6 
7 std::unique_ptr<optimization_guide::proto::PredictionModel>
GetMinimalDecisionTreePredictionModel(double threshold,double weight)8 GetMinimalDecisionTreePredictionModel(double threshold, double weight) {
9   auto prediction_model =
10       std::make_unique<optimization_guide::proto::PredictionModel>();
11   prediction_model->mutable_model()->mutable_threshold()->set_value(threshold);
12   optimization_guide::proto::DecisionTree decision_tree_model;
13   decision_tree_model.set_weight(weight);
14 
15   optimization_guide::proto::TreeNode* tree_node =
16       decision_tree_model.add_nodes();
17   tree_node->mutable_node_id()->set_value(0);
18 
19   *prediction_model->mutable_model()->mutable_decision_tree() =
20       decision_tree_model;
21 
22   return prediction_model;
23 }
24 
25 std::unique_ptr<optimization_guide::proto::PredictionModel>
GetSingleLeafDecisionTreePredictionModel(double threshold,double weight,double leaf_value)26 GetSingleLeafDecisionTreePredictionModel(double threshold,
27                                          double weight,
28                                          double leaf_value) {
29   auto prediction_model =
30       GetMinimalDecisionTreePredictionModel(threshold, weight);
31   prediction_model->mutable_model()
32       ->mutable_decision_tree()
33       ->mutable_nodes(0)
34       ->mutable_leaf()
35       ->mutable_vector()
36       ->add_value()
37       ->set_double_value(leaf_value);
38   return prediction_model;
39 }
40