Add test for invalid version

keras-team · Sep 19, 2023 · 0e7cbb9 · 0e7cbb9
1 parent f53cfd1
commit 0e7cbb9
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/keras_core/export/export_lib_test.py b/keras_core/export/export_lib_test.py
@@ -542,6 +542,18 @@ def test_model_export_method(self):
         )
 
 
+@pytest.mark.skipif(
+    backend.backend() != "jax" or sys.modules["jax"].__version__ <= "0.4.15",
+    reason="This test is for invalid JAX versions, i.e. versions > 0.4.15.",
+)
+class VersionTest(testing.TestCase):
+    def test_invalid_jax_version(self):
+        with self.assertRaisesRegex(
+            ValueError, "only compatible with JAX version"
+        ):
+            _ = export_lib.ExportArchive()
+
+
 @pytest.mark.skipif(
     backend.backend() != "tensorflow",
     reason="TFSM Layer reloading is only for the TF backend.",