xref: /aosp_15_r20/external/pytorch/test/test_utils_internal.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# Owner(s): ["module: unknown"]
2
3import os
4
5from torch._utils_internal import justknobs_feature, JustKnobsConfig
6from torch.testing._internal.common_utils import (  # type: ignore[attr-defined]
7    load_tests,
8)
9
10
11# load_tests from torch.testing._internal.common_utils is used to automatically filter tests for
12# sharding on sandcastle. This line silences flake warnings
13load_tests = load_tests
14
15from torch.testing._internal.common_utils import run_tests, TestCase
16
17
18class TestJustKnob(TestCase):
19    def test_justknob_config(self):
20        with self.subTest("Returns True"):
21            a = JustKnobsConfig()
22            self.assertTrue(a.get())
23        with self.subTest("Returns False"):
24            a = JustKnobsConfig(name="fake_name", default=False)
25            self.assertFalse(a.get())
26        with self.subTest("Returns True via config"):
27            a = JustKnobsConfig(name="fake_name", default=False)
28            a.set(True)
29            self.assertTrue(a.get())
30        with self.subTest("Returns True via env"):
31            os.environ["FAKE_FEATURE"] = "1"
32            a = JustKnobsConfig(
33                name="fake_name", env_name="FAKE_FEATURE", default=False
34            )
35            self.assertTrue(a.get())
36        with self.subTest("Returns same value consistently"):
37            a = JustKnobsConfig(name="fake_name", default=False)
38            a.set(True)
39            self.assertTrue(a.get())
40            a.set(False)
41            self.assertTrue(a.get())
42        with self.subTest("Checks __bool__"):
43            a = JustKnobsConfig(name="fake_name", default=False)
44            if a:
45                raise RuntimeError("Should not be true")
46            self.assertFalse(a)
47
48    def test_justknob_feature(self):
49        with self.subTest("OSS is True"):
50            self.assertTrue(justknobs_feature("testname"))
51        with self.subTest("OSS default=True"):
52            self.assertTrue(justknobs_feature("testname", default=True))
53        with self.subTest("OSS default=False"):
54            self.assertFalse(justknobs_feature("testname", default=False))
55        with self.subTest("OSS config=True, default=False"):
56            self.assertTrue(
57                justknobs_feature("testname", config_value=True, default=False)
58            )
59        with self.subTest("OSS config=None, default=False"):
60            self.assertFalse(
61                justknobs_feature("testname", config_value=None, default=False)
62            )
63        with self.subTest("OSS config=False, default=True"):
64            self.assertFalse(
65                justknobs_feature("testname", config_value=False, default=True)
66            )
67        with self.subTest("OSS env is missing, config=False, default=True"):
68            self.assertFalse(
69                justknobs_feature(
70                    "testname", config_value=False, env_name="NOTDEFINED", default=False
71                )
72            )
73        with self.subTest("OSS env is missing, default=False"):
74            self.assertFalse(
75                justknobs_feature("testname", env_name="NOTDEFINED", default=False)
76            )
77        with self.subTest(
78            "OSS config overrides env, config=True, env=False, default=False"
79        ):
80            os.environ["FEATURE_ENV"] = "0"
81            self.assertTrue(
82                justknobs_feature(
83                    "testname",
84                    config_value=True,
85                    env_name="FEATURE_ENV",
86                    default=False,
87                )
88            )
89        with self.subTest("OSS env overrides default, , default=False"):
90            os.environ["FEATURE_ENV"] = "1"
91            self.assertTrue(
92                justknobs_feature("testname", env_name="FEATURE_ENV", default=False)
93            )
94        with self.subTest("OSS env truthy, config=False, default=False"):
95            os.environ["FEATURE_ENV"] = "1"
96            self.assertTrue(
97                justknobs_feature(
98                    "testname",
99                    env_name="FEATURE_ENV",
100                    default=False,
101                )
102            )
103            os.environ["FEATURE_ENV"] = "true"
104            self.assertTrue(
105                justknobs_feature(
106                    "testname",
107                    env_name="FEATURE_ENV",
108                    default=False,
109                )
110            )
111            os.environ["FEATURE_ENV"] = "TRUE"
112            self.assertTrue(
113                justknobs_feature(
114                    "testname",
115                    env_name="FEATURE_ENV",
116                    default=False,
117                )
118            )
119            os.environ["FEATURE_ENV"] = "very weird true"
120            self.assertTrue(
121                justknobs_feature(
122                    "testname",
123                    env_name="FEATURE_ENV",
124                    default=False,
125                )
126            )
127        with self.subTest("OSS env false, default=True"):
128            os.environ["FEATURE_ENV"] = "0"
129            self.assertFalse(
130                justknobs_feature("testname", env_name="FEATURE_ENV", default=True)
131            )
132            os.environ["FEATURE_ENV"] = "false"
133            self.assertFalse(
134                justknobs_feature("testname", env_name="FEATURE_ENV", default=True)
135            )
136            os.environ["FEATURE_ENV"] = "FALSE"
137            self.assertFalse(
138                justknobs_feature("testname", env_name="FEATURE_ENV", default=True)
139            )
140
141
142if __name__ == "__main__":
143    run_tests()
144