Merge pull request #3 from google/test-tidy

Test tidy

test/conftest.py

@@ -27,7 +27,11 @@ import typeguard
 try:
     import beartype
 except ImportError:
-    typecheck_params = [typeguard.typechecked]
+
+    def skip(*args, **kwargs):
+        pytest.skip("Beartype not installed")
+
+    typecheck_params = [typeguard.typechecked, skip]
 else:
     typecheck_params = [typeguard.typechecked, beartype.beartype]
 

test/import_hook_tester_beartype.py

@@ -19,10 +19,11 @@
 
 import jax.numpy as jnp
 import pytest
-from helpers import ParamError
 
 from jaxtyping import f32
 
+from .helpers import ParamError
+
 
 def g(x: f32[" b"]):
     pass

test/import_hook_tester_broken_checker.py

@@ -19,10 +19,11 @@
 
 import jax.numpy as jnp
 import pytest
-from helpers import ParamError
 
 from jaxtyping import f32
 
+from .helpers import ParamError
+
 
 def g(x: f32[" b"]):
     pass

test/import_hook_tester_transitive/another_file.py

@@ -19,10 +19,11 @@
 
 import jax.numpy as jnp
 import pytest
-from helpers import ParamError
 
 from jaxtyping import f32
 
+from ..helpers import ParamError
+
 
 def g(x: f32[" b"]):
     pass

test/import_hook_tester_typeguard.py

@@ -19,10 +19,11 @@
 
 import jax.numpy as jnp
 import pytest
-from helpers import ParamError
 
 from jaxtyping import f32
 
+from .helpers import ParamError
+
 
 def g(x: f32[" b"]):
     pass

test/test_array.py

@@ -20,10 +20,11 @@
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
-from helpers import ParamError, ReturnError
 
 from jaxtyping import Array, f, f32, jaxtyped
 
+from .helpers import ParamError, ReturnError
+
 
 def test_basic(typecheck):
     @jaxtyped

test/test_import_hook.py

@@ -24,9 +24,9 @@ from jaxtyping import install_import_hook
 
 def test_import_hook_typeguard():
     hook = install_import_hook(
-        "import_hook_tester_typeguard", ("typeguard", "typechecked")
+        "test.import_hook_tester_typeguard", ("typeguard", "typechecked")
     )
-    import import_hook_tester_typeguard  # noqa: F401
+    from . import import_hook_tester_typeguard  # noqa: F401
 
     hook.uninstall()
 
@@ -35,29 +35,29 @@ def test_import_hook_beartype():
     try:
         import beartype  # noqa: F401
     except ImportError:
-        pass
+        pytest.skip("Beartype not installed")
     else:
         hook = install_import_hook(
-            "import_hook_tester_beartype", ("beartype", "beartype")
+            "test.import_hook_tester_beartype", ("beartype", "beartype")
         )
-        import import_hook_tester_beartype  # noqa: F401
+        from . import import_hook_tester_beartype  # noqa: F401
 
         hook.uninstall()
 
 
 def test_import_hook_transitive():
     hook = install_import_hook(
-        "import_hook_tester_transitive", ("typeguard", "typechecked")
+        "test.import_hook_tester_transitive", ("typeguard", "typechecked")
     )
-    import import_hook_tester_transitive  # noqa: F401
+    from . import import_hook_tester_transitive  # noqa: F401
 
     hook.uninstall()
 
 
 def test_import_hook_broken_checker():
     hook = install_import_hook(
-        "import_hook_tester_broken_checker", ("jaxtyping", "does_not_exist")
+        "test.import_hook_tester_broken_checker", ("jaxtyping", "does_not_exist")
     )
     with pytest.raises(AttributeError):
-        import import_hook_tester_broken_checker  # noqa: F401
+        from . import import_hook_tester_broken_checker  # noqa: F401
     hook.uninstall()

test/test_pytree.py

@@ -24,10 +24,11 @@ import jax
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
-from helpers import make_mlp, ParamError
 
 from jaxtyping import f, jaxtyped, PyTree
 
+from .helpers import make_mlp, ParamError
+
 
 def test_direct(typecheck):
     @typecheck