# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0
# This file is for testing ONNX with ONNX Runtime
# Create a general scenario to use ONNX Runtime with ONNX

import unittest


class TestONNXRuntime(unittest.TestCase):
    def test_with_ort_example(self) -> None:
        try:
            import onnxruntime

            del onnxruntime
        except ImportError:
            raise unittest.SkipTest("onnxruntime not installed") from None

        from numpy import float32, random
        from onnxruntime import InferenceSession
        from onnxruntime.datasets import get_example

        from onnx import checker, load, shape_inference, version_converter

        # get certain example model from ORT using opset 9
        example1 = get_example("sigmoid.onnx")

        # test ONNX functions
        model = load(example1)
        checker.check_model(model)
        checker.check_model(model, full_check=True)
        inferred_model = shape_inference.infer_shapes(
            model, check_type=True, strict_mode=True, data_prop=True
        )
        converted_model = version_converter.convert_version(inferred_model, 10)

        # test ONNX Runtime functions
        sess = InferenceSession(
            converted_model.SerializeToString(), providers=["CPUExecutionProvider"]
        )
        input_name = sess.get_inputs()[0].name
        output_name = sess.get_outputs()[0].name
        x = random.random((3, 4, 5))
        x = x.astype(float32)

        sess.run([output_name], {input_name: x})


if __name__ == "__main__":
    unittest.main()