import unittest
import torch
from tensorboardX import SummaryWriter


class ONNXGraphTest(unittest.TestCase):
    def test_onnx_graph(self):
        import subprocess
        zoo_address = 'https://onnxzoo.blob.core.windows.net/models/opset_8/mnist/mnist.tar.gz'

        res = subprocess.call(['wget', '-nc', zoo_address])
        assert res == 0, 'cannot download example onnx model from the zoo'
        res = subprocess.call(['tar', 'xf', 'mnist.tar.gz', '-C', 'examples/', 'mnist/model.onnx'])

        with SummaryWriter() as w:
            w.add_onnx_graph('examples/mnist/model.onnx')