@@ -0,0 +1,2776 @@
+#%0A#  Copyright 2019 The FATE Authors. All Rights Reserved.%0A#%0A#  Licensed under the Apache License, Version 2.0 (the %22License%22);%0A#  you may not use this file except in compliance with the License.%0A#  You may obtain a copy of the License at%0A#%0A#      http://www.apache.org/licenses/LICENSE-2.0%0A#%0A#  Unless required by applicable law or agreed to in writing, software%0A#  distributed under the License is distributed on an %22AS IS%22 BASIS,%0A#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.%0A#  See the License for the specific language governing permissions and%0A#  limitations under the License.%0A#%0A%0Aimport numpy as np%0Aimport unittest%0Aimport uuid%0A%0Afrom arch.api import session%0Afrom federatedml.feature.instance import Instance%0Afrom federatedml.local_baseline.local_baseline import LocalBaseline%0Afrom sklearn.linear_model import LogisticRegression%0A%0Aclass TestLocalBaselin(unittest.TestCase):%0A    def setUp(self):%0A        self.job_id = str(uuid.uuid1())%0A        session.init(self.job_id)%0A        data_num = 100%0A        feature_num = 8%0A        self.prepare_data(data_num, feature_num)%0A        local_baseline_obj = LocalBaseline()%0A        local_baseline_obj.need_run = True%0A        local_baseline_obj.header = %5B%22x1%22, %22x2%22, %22x3%22, %22x4%22, %22x5%22, %22x6%22, %22x7%22, %22x8%22%5D%0A        local_baseline_obj.model_name = %22LogisticRegression%22%0A        local_baseline_obj.model_opts = %7B%7D%0A        self.local_baseline_obj = local_baseline_obj%0A%0A    def prepare_data(self, data_num, feature_num):%0A        self.X = np.random.randint(0, 10, (data_num, feature_num))%0A        self.y = np.random.randint(0, 2, data_num)%0A        final_result = %5B%5D%0A        for i in range(data_num):%0A            tmp = self.X%5Bi,:%5D%0A            inst = Instance(inst_id=i, features=tmp, label=self.y%5Bi%5D)%0A            tmp = (str(i), inst)%0A            final_result.append(tmp)%0A        table = session.parallelize(final_result,%0A                                    include_key=True,%0A                                    partition=3)%0A        self.table = table%0A%0A    def test_predict(self):%0A        glm = LogisticRegression().fit(self.X, self.y)%0A        real_predict_result = glm.predict(self.X)%0A        self.local_baseline_obj.model_fit = glm%0A        model_predict_result = self.local_baseline_obj.predict(self.table)%0A        model_predict_result = np.array(%5Bv%5B1%5D%5B1%5D for v in model_predict_result.collect()%5D)%0A%0A        np.testing.assert_array_equal(model_predict_result, real_predict_result)%0A%0A    def tearDown(self):%0A         session.stop()%0A         try:%0A             session.cleanup(%22*%22, self.job_id, True)%0A         except EnvironmentError:%0A             pass%0A         try:%0A             session.cleanup(%22*%22, self.job_id, False)%0A         except EnvironmentError:%0A             pass%0A%0A%0Aif __name__ == '__main__':%0A    unittest.main()%0A

add local baseline test