1# Owner(s): ["oncall: mobile"] 2 3import os 4import tempfile 5from pathlib import Path 6 7from torch.jit.generate_bytecode import generate_upgraders_bytecode 8from torch.testing._internal.common_utils import run_tests, TestCase 9from torchgen.operator_versions.gen_mobile_upgraders import sort_upgrader, write_cpp 10 11 12pytorch_caffe2_dir = Path(__file__).resolve().parents[2] 13 14 15class TestLiteScriptModule(TestCase): 16 def test_generate_bytecode(self): 17 upgrader_list = generate_upgraders_bytecode() 18 sorted_upgrader_list = sort_upgrader(upgrader_list) 19 upgrader_mobile_cpp_path = ( 20 pytorch_caffe2_dir 21 / "torch" 22 / "csrc" 23 / "jit" 24 / "mobile" 25 / "upgrader_mobile.cpp" 26 ) 27 with tempfile.TemporaryDirectory() as tmpdirname: 28 write_cpp(tmpdirname, sorted_upgrader_list) 29 with open(os.path.join(tmpdirname, "upgrader_mobile.cpp")) as file_name: 30 actual_output = [line.strip() for line in file_name if line] 31 with open(str(upgrader_mobile_cpp_path)) as file_name: 32 expect_output = [line.strip() for line in file_name if line] 33 actual_output_filtered = list( 34 filter(lambda token: len(token) != 0, actual_output) 35 ) 36 expect_output_filtered = list( 37 filter(lambda token: len(token) != 0, expect_output) 38 ) 39 40 self.assertEqual(actual_output_filtered, expect_output_filtered) 41 42 43if __name__ == "__main__": 44 run_tests() 45