diff --git a/ethosu/mlw_codec/test/test_mlw_codec.py b/ethosu/mlw_codec/test/test_mlw_codec.py
new file mode 100644
index 0000000..31a3bc0
--- /dev/null
+++ b/ethosu/mlw_codec/test/test_mlw_codec.py
@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+# Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the License); you may
+# not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an AS IS BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Simple example of the usage of mlw_codec.
+import pytest
+from ethosu import mlw_codec
+
+
+class TestMLWCodec:
+    """ This class is responsible to test the mlw_codec library
+    It mainly tests the two methods encode() and decode() with different inputs"""
+
+    weights = [0, 2, 3, 0, -1, -2, -3, 0, 0, 0, 1, -250, 240] * 3
+    compressed_weights = bytearray(
+        b"\xb8\x00\\q^\x1f\xfc\x01\x03\x05\x08\x0c\x10\x908\x12\xd7\x99:\xd2\x99$\xae#\x9d\xa9#\x00\xf0\xff\xff\xff"
+    )
+    empty_decoded = bytearray(b"\xfe\xffC\x00\xf0\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff")
+
+    # Generate parameters lists for the tests below
+    encode_testdata = [
+        (mlw_codec.encode, weights, compressed_weights),
+        pytest.param(mlw_codec.encode, ["a"], empty_decoded, marks=pytest.mark.xfail),  # cannot accept strings
+    ]
+
+    decode_testdata = [(mlw_codec.decode, compressed_weights, weights)]
+
+    codec_testdata = [
+        (weights, weights),
+        ([1] * 10, [1] * 10),
+        pytest.param(["a"], ["a"], marks=pytest.mark.xfail),  # cannot accept strings
+    ]
+
+    @pytest.mark.parametrize("function_under_test,test_input,expected", encode_testdata)
+    def test_mlw_codec(self, function_under_test, test_input, expected):
+        self._call_mlw_codec_method(function_under_test, test_input, expected)
+
+    @pytest.mark.parametrize("function_under_test,test_input,expected", decode_testdata)
+    def test_mlw_decode(self, function_under_test, test_input, expected):
+        self._call_mlw_codec_method(function_under_test, test_input, expected)
+
+    @pytest.mark.parametrize("test_input,expected", codec_testdata)
+    def test_mlw_encode_decode(self, test_input, expected):
+        output = mlw_codec.decode(mlw_codec.encode(test_input))
+        assert output == expected
+
+    def _call_mlw_codec_method(self, method_name, test_input, expected):
+        output = method_name(test_input)
+        assert output == expected
diff --git a/ethosu/mlw_codec/test_mlw_codec.py b/ethosu/mlw_codec/test_mlw_codec.py
deleted file mode 100644
index 0bcd417..0000000
--- a/ethosu/mlw_codec/test_mlw_codec.py
+++ /dev/null
@@ -1,38 +0,0 @@
-#!/usr/bin/env python3
-# Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
-#
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the License); you may
-# not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an AS IS BASIS, WITHOUT
-# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# Simple example of the usage of mlw_codec.
-import sys
-
-from ethosu import mlw_codec
-
-# Simple example
-if __name__ == "__main__":
-    weights = [0, 2, 3, 0, -1, -2, -3, 0, 0, 0, 1, -250, 240] * 3
-    print("Original weights    :", weights)
-
-    compressed_weights = mlw_codec.encode(weights)
-    print("Compressed weights  :", len(compressed_weights), compressed_weights)
-
-    uncompressed_weights = mlw_codec.decode(compressed_weights)
-    print("Uncompressed weights:", uncompressed_weights)
-
-    if weights != uncompressed_weights:
-        print("TEST FAILED")
-        sys.exit(1)
-    else:
-        print("TEST PASSED")
-        sys.exit(0)