Skip to content

Commit

Permalink
Move externalization tests in test file
Browse files Browse the repository at this point in the history
  • Loading branch information
vinayakdsci committed Nov 14, 2024
1 parent 0836529 commit c586034
Showing 1 changed file with 24 additions and 24 deletions.
48 changes: 24 additions & 24 deletions compiler/bindings/python/test/tools/import_onnx_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,30 @@ def run_tool(*argv: str):
)


class ImportOnnxTest(unittest.TestCase):
def setUp(self):
with tempfile.NamedTemporaryFile(delete=False) as f:
self.outputPath = f.name

def tearDown(self) -> None:
if os.path.exists(self.outputPath):
os.unlink(self.outputPath)

def testConsoleOutput(self):
# Just test that it doesn't crash: rely on the file test for verification.
run_tool(ONNX_FILE_PATH)

def testDisableVerify(self):
# Just test that the flag is accepted.
run_tool(ONNX_FILE_PATH, "--no-verify")

def testFileOutput(self):
run_tool(ONNX_FILE_PATH, "-o", self.outputPath)
with open(self.outputPath, "rt") as f:
contents = f.read()
self.assertIn("torch.operator", contents)


class ImportOnnxwithExternalizationTest(unittest.TestCase):
def setUp(self):
with tempfile.NamedTemporaryFile(delete=False) as f:
Expand Down Expand Up @@ -102,30 +126,6 @@ def testExternalizeMinimumThreshold(self):
self.assertNotIn("onnx.Constant", contents)


class ImportOnnxTest(unittest.TestCase):
def setUp(self):
with tempfile.NamedTemporaryFile(delete=False) as f:
self.outputPath = f.name

def tearDown(self) -> None:
if os.path.exists(self.outputPath):
os.unlink(self.outputPath)

def testConsoleOutput(self):
# Just test that it doesn't crash: rely on the file test for verification.
run_tool(ONNX_FILE_PATH)

def testDisableVerify(self):
# Just test that the flag is accepted.
run_tool(ONNX_FILE_PATH, "--no-verify")

def testFileOutput(self):
run_tool(ONNX_FILE_PATH, "-o", self.outputPath)
with open(self.outputPath, "rt") as f:
contents = f.read()
self.assertIn("torch.operator", contents)


if __name__ == "__main__":
try:
import onnx
Expand Down

0 comments on commit c586034

Please sign in to comment.