From 72d863f6fdb4f8606555625ce43ed00d58f6f088 Mon Sep 17 00:00:00 2001 From: Harsimran Singh Dalal Date: Mon, 28 Oct 2024 21:04:20 +0530 Subject: [PATCH] Create test_model_loading.py --- test_model_loading.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 test_model_loading.py diff --git a/test_model_loading.py b/test_model_loading.py new file mode 100644 index 0000000..0994028 --- /dev/null +++ b/test_model_loading.py @@ -0,0 +1,18 @@ +import unittest +from model_optimization import load_optimized_model # Replace with the actual module name + +class TestModelLoading(unittest.TestCase): + def test_model_loading(self): + """Test if the optimized model loads correctly.""" + model = load_optimized_model() + self.assertIsNotNone(model, "The model should not be None after loading.") + self.assertTrue(hasattr(model, 'predict'), "The loaded model should have a 'predict' method.") + + def test_model_input_shape(self): + """Test if the model input shape is as expected.""" + model = load_optimized_model() + input_shape = model.input_shape + self.assertEqual(input_shape, (None, 224, 224, 3), "Model input shape should be (None, 224, 224, 3).") + +if __name__ == "__main__": + unittest.main()