ml_pid_cbm.tools.test_json_tools

 1import unittest
 2from unittest.mock import mock_open, patch
 3
 4from . import json_tools
 5
 6
 7class TestJsonTools(unittest.TestCase):
 8    @classmethod
 9    def setUpClass(cls):
10        cls.json_data = """{"file_names":{
11        "training": "training_filename.test",
12        "test": "test_filename.test"
13        },"var_names":{"momentum": "Complex_p","charge": "Complex_q"},
14        "cuts":{"Complex_mass2": {"lower": -1.0,"upper": 2.0},"Complex_pT": {"lower": 0.0,"upper": 2.0}},
15        "features_for_train":["Complex_mass2", "Complex_p"],
16        "vars_to_draw":["Complex_mass2", "Complex_eta"],
17        "hyper_params": {"values": {"n_estimators": 596,"max_depth": 5,"learning_rate": 0.07161792803939408},
18        "ranges": {"n_estimators": [400, 1000],"max_depth": [2, 6],"learning_rate": [0.01, 0.1]}}}"""
19
20    def test_create_cut_string(self):
21        # ideally formatted
22        expected_string = "0.1 <= test_cut < 13.0"
23        self.assertEqual(
24            json_tools.create_cut_string(0.1, 13.0, "test_cut"), expected_string
25        )
26        # not ideally formatted
27        self.assertEqual(
28            json_tools.create_cut_string(0.141, 13, "test_cut"), expected_string
29        )
30
31    def test_load_var_name(self):
32        # mocking json file for testing
33        with patch("builtins.open", mock_open(read_data=self.json_data)):
34            self.assertEqual(
35                json_tools.load_var_name("test.json", "momentum"), "Complex_p"
36            )
37
38    def test_load_quality_cuts(self):
39        # mocking json file for testing
40        with patch("builtins.open", mock_open(read_data=self.json_data)):
41            quality_cuts = json_tools.load_quality_cuts("test.json")
42            expected_cuts = ["-1.0 <= Complex_mass2 < 2.0", "0.0 <= Complex_pT < 2.0"]
43            self.assertEqual(quality_cuts, expected_cuts)
44
45    def test_load_features_for_train(self):
46        target_features = ["Complex_mass2", "Complex_p"]
47        with patch("builtins.open", mock_open(read_data=self.json_data)):
48            read_features = json_tools.load_features_for_train("test.json")
49            self.assertEqual(target_features, read_features)
50    def test_load_vars_to_draw(self):
51        target_vars = ["Complex_mass2", "Complex_eta"]
52        with patch("builtins.open", mock_open(read_data=self.json_data)):
53            read_vars = json_tools.load_vars_to_draw("test.json")
54            self.assertEqual(target_vars, read_vars)
55
56    def test_load_hyper_params_vals(self):
57        target_vals = (596, 5, 0.07161792803939408)
58        with patch("builtins.open", mock_open(read_data=self.json_data)):
59            read_vals = json_tools.load_hyper_params_vals("test.json")
60            self.assertEqual(target_vals, read_vals)
61
62    def test_load_file_name(self):
63            target_name = "training_filename.test"
64            with patch("builtins.open", mock_open(read_data=self.json_data)):
65                read_name = json_tools.load_file_name("test.json", training_or_test="training")
66                self.assertEqual(target_name, read_name)
67