# Copyright 2019 ZTE corporation. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

from os import path
from tempfile import NamedTemporaryFile, TemporaryDirectory
from unittest import TestCase

import tensorflow as tf
from tensorflow import keras

import model_compiler


class CompileModelTestCase(TestCase):
    def test_error_result(self):
        result = model_compiler.compile_model({'serving_type': 'foo',
                                               'model_name': 'bar',
                                               'max_batch_size': 4,
                                               'export_path': '/foo/bar'})

        self.assertEqual(result, {'status': 'failure', 'error_msg': 'Unable to determine the source model type.'})

    def test_success_result(self):
        with NamedTemporaryFile(suffix='.h5') as model_file, TemporaryDirectory() as target_dir:
            with tf.compat.v1.Session(graph=tf.Graph()):
                keras.Sequential([keras.layers.Dense(units=4, input_shape=[8])]).save(model_file.name)

            result = model_compiler.compile_model({
                'serving_type': 'tf',
                'model_name': 'foobar',
                'version': 4,
                'max_batch_size': 7,
                'input_model': model_file.name,
                'input_signatures': ['x'],
                'output_signatures': ['y'],
                'export_path': target_dir
            })

            self.assertEqual(result, {'status': 'success', 'path': path.join(target_dir, 'foobar_4.zip')})
