pytorch

1
# Owner(s): ["module: unknown"]
2

3
import os
4

5
from torch._utils_internal import justknobs_feature, JustKnobsConfig
6
from 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
13
load_tests = load_tests
14

15
from torch.testing._internal.common_utils import run_tests, TestCase
16

17

18
class 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

142
if __name__ == "__main__":
143
    run_tests()
144