TensorFlowAOT/test/testAOTTools.py

0001 # coding: utf-8
0002
0003 """
0004 AOT compilation and dev workflow tests.
0005 """
0006
0007 import os
0008 import re
0009 import shlex
0010 import subprocess
0011 import tempfile
0012 import functools
0013 import unittest
0014
0015
0016 this_dir = os.path.dirname(os.path.abspath(__file__))
0017
0018
0019 def run_cmd(cmd):
0020     if not isinstance(cmd, str):
0021         cmd = shlex.join(cmd)
0022     return subprocess.run(cmd, shell=True, check=True)
0023
0024
0025 def run_in_tmp(func):
0026     @functools.wraps(func)
0027     def wrapper(self):
0028         with tempfile.TemporaryDirectory() as tmp_dir:
0029             func(self, tmp_dir)
0030     return wrapper
0031
0032
0033 class TFAOTTests(unittest.TestCase):
0034
0035     @run_in_tmp
0036     def test_dev_workflow(self, tmp_dir):
0037         import cms_tfaot
0038
0039         # find the cms_tfaot install dir to locate the test model
0040         m = re.match(r"(.+/\d+\.\d+\.\d+\-[^/]+)/lib/.+$", cms_tfaot.__file__)
0041         self.assertIsNotNone(m)
0042         config_file = os.path.join(m.group(1), "share", "test_models", "simple", "aot_config.yaml")
0043         self.assertTrue(os.path.exists(config_file))
0044
0045         # run the dev workflow
0046         # create the test model
0047         cmd = [
0048             "cms_tfaot_compile",
0049             "-c", config_file,
0050             "-o", tmp_dir,
0051             "--tool-name", "tfaot-model-test",
0052             "--dev",
0053         ]
0054         run_cmd(cmd)
0055
0056         # check files
0057         exists = lambda *p: os.path.exists(os.path.join(tmp_dir, *p))
0058         self.assertTrue(exists("tfaot-model-test.xml"))
0059         self.assertTrue(exists("include", "tfaot-model-test"))
0060         self.assertTrue(exists("include", "tfaot-model-test", "test_simple_bs1.h"))
0061         self.assertTrue(exists("include", "tfaot-model-test", "test_simple_bs2.h"))
0062         self.assertTrue(exists("include", "tfaot-model-test", "model.h"))
0063         self.assertTrue(exists("lib", "test_simple_bs1.o"))
0064         self.assertTrue(exists("lib", "test_simple_bs2.o"))
0065
0066
0067 if __name__ == "__main__":
0068     unittest.main()