1# Lint as: python3 2# Copyright 2019 The TensorFlow Authors. All Rights Reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15# ============================================================================== 16 17"""Test for data_split_person.py.""" 18 19from __future__ import absolute_import 20from __future__ import division 21from __future__ import print_function 22 23import unittest 24from data_split_person import person_split 25from data_split_person import read_data 26 27 28class TestSplitPerson(unittest.TestCase): 29 30 def setUp(self): # pylint: disable=g-missing-super-call 31 self.data = read_data("./data/complete_data") 32 33 def test_person_split(self): 34 train_names = ["dengyl"] 35 valid_names = ["liucx"] 36 test_names = ["tangsy"] 37 dengyl_num = 63 38 liucx_num = 63 39 tangsy_num = 30 40 train_data, valid_data, test_data = person_split(self.data, train_names, 41 valid_names, test_names) 42 self.assertEqual(len(train_data), dengyl_num) 43 self.assertEqual(len(valid_data), liucx_num) 44 self.assertEqual(len(test_data), tangsy_num) 45 self.assertIsInstance(train_data, list) 46 self.assertIsInstance(valid_data, list) 47 self.assertIsInstance(test_data, list) 48 self.assertIsInstance(train_data[0], dict) 49 self.assertIsInstance(valid_data[0], dict) 50 self.assertIsInstance(test_data[0], dict) 51 52 53if __name__ == "__main__": 54 unittest.main() 55