Test cleanup and formatting issues addressed

Signed-off-by: Christopher Horvath <chorvath@nvidia.com>

Author: blackencino

point_transformer_v3/tests/conftest.py

@@ -13,14 +13,16 @@
 import sys
 from unittest.mock import MagicMock
 
+
 # Helper to mock a module and its submodules
 def mock_module(module_name):
     if module_name not in sys.modules:
         m = MagicMock()
-        m.__path__ = [] # Make it look like a package
+        m.__path__ = []  # Make it look like a package
         sys.modules[module_name] = m
     return sys.modules[module_name]
 
+
 # Mock pointops to avoid installing custom CUDA extensions for unit tests
 mock_module("pointops")
 mock_module("pointgroup_ops")

point_transformer_v3/tests/unit/__init__.py

@@ -1,3 +1,2 @@
 # Copyright Contributors to the OpenVDB Project
 # SPDX-License-Identifier: Apache-2.0
-

point_transformer_v3/tests/unit/test_order_type_validation.py

@@ -2,90 +2,129 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Unit tests for order type validation in PTV3_Attention.
+Unit tests for space-filling curve functions in fvdb_utils.
 
-This is a minimal unit test to establish the test infrastructure.
-Tests a pure function with no dependencies on FVDB operations or external libraries.
+Tests the standalone curve type validation and dispatch logic
+in space_filling_curve_from_jagged_ijk and related functions.
 """
 
 import unittest
-
-from fvdb_extensions.models.ptv3_fvdb import PTV3_Attention
-
-
-class TestOrderTypeValidation(unittest.TestCase):
-    """Test cases for order type validation in PTV3_Attention."""
-
-    def setUp(self):
-        """Set up test fixtures."""
-        pass
-
-    def test_permute_valid_order_types(self):
-        """Test that _permute accepts valid order types."""
-        # Create a minimal attention instance (we won't call forward, just test _permute)
-        attn = PTV3_Attention(
-            hidden_size=64,
-            num_heads=1,
-            proj_drop=0.0,
-            patch_size=0,
-            order_index=0,
-            order_types=("z", "z-trans"),
-        )
-
-        # Test that valid order types don't raise errors
-        # Note: We can't actually call _permute without a grid, but we can test the validation logic
-        valid_types = ["z", "z-trans", "hilbert", "hilbert-trans"]
-        for order_type in valid_types:
-            # The _permute method should accept these without raising ValueError
-            # We test this by checking the method exists and has the right signature
-            self.assertTrue(hasattr(attn, "_permute"), "Should have _permute method")
-            self.assertTrue(callable(getattr(attn, "_permute")), "_permute should be callable")
-
-    def test_permute_invalid_order_type_raises(self):
-        """Test that _permute raises ValueError for invalid order types."""
-        attn = PTV3_Attention(
-            hidden_size=64,
-            num_heads=1,
-            proj_drop=0.0,
-            patch_size=0,
-            order_index=0,
-            order_types=("z",),
-        )
-
-        # Create a mock grid object with minimal interface needed for _permute
-        class MockGrid:
-            def morton(self):
-                return None
-
-            def morton_zyx(self):
-                return None
-
-            def hilbert(self):
-                return None
-
-            def hilbert_zyx(self):
-                return None
-
-        mock_grid = MockGrid()
-
-        # Test that invalid order type raises ValueError
-        with self.assertRaises(ValueError):
-            attn._permute(mock_grid, "invalid_order_type")
-
-    def test_order_type_initialization(self):
-        """Test that order_types are correctly stored during initialization."""
-        order_types = ("z", "z-trans", "hilbert")
-        attn = PTV3_Attention(
-            hidden_size=64,
-            num_heads=1,
-            proj_drop=0.0,
-            patch_size=0,
-            order_index=0,
-            order_types=order_types,
-        )
-        self.assertEqual(attn.order_types, order_types, "Order types should be stored correctly")
+from unittest.mock import MagicMock, patch
+
+from fvdb_extensions.models.fvdb_utils import (
+    hilbert_flipped_from_jagged_ijk,
+    hilbert_from_jagged_ijk,
+    identity_from_jagged_ijk,
+    morton_flipped_from_jagged_ijk,
+    morton_from_jagged_ijk,
+    space_filling_curve_from_jagged_ijk,
+)
+
+
+class TestSpaceFillingCurveValidation(unittest.TestCase):
+    """Test cases for curve type validation in space_filling_curve_from_jagged_ijk."""
+
+    def test_invalid_curve_type_raises_value_error(self):
+        """Test that invalid curve types raise ValueError."""
+        mock_jagged_ijk = MagicMock()
+
+        invalid_types = ["invalid", "unknown", "random", "xyz", ""]
+        for curve_type in invalid_types:
+            with self.assertRaises(ValueError, msg=f"Should raise ValueError for '{curve_type}'"):
+                space_filling_curve_from_jagged_ijk(mock_jagged_ijk, curve_type)
+
+    @patch("fvdb_extensions.models.fvdb_utils.morton_from_jagged_ijk")
+    def test_morton_alias_z(self, mock_morton):
+        """Test that 'z' dispatches to morton_from_jagged_ijk."""
+        mock_jagged_ijk = MagicMock()
+        mock_morton.return_value = MagicMock()
+
+        space_filling_curve_from_jagged_ijk(mock_jagged_ijk, "z")
+        mock_morton.assert_called_once_with(mock_jagged_ijk)
+
+    @patch("fvdb_extensions.models.fvdb_utils.morton_from_jagged_ijk")
+    def test_morton_alias_morton(self, mock_morton):
+        """Test that 'morton' dispatches to morton_from_jagged_ijk."""
+        mock_jagged_ijk = MagicMock()
+        mock_morton.return_value = MagicMock()
+
+        space_filling_curve_from_jagged_ijk(mock_jagged_ijk, "morton")
+        mock_morton.assert_called_once_with(mock_jagged_ijk)
+
+    @patch("fvdb_extensions.models.fvdb_utils.morton_flipped_from_jagged_ijk")
+    def test_morton_flipped_alias_z_trans(self, mock_morton_flipped):
+        """Test that 'z-trans' dispatches to morton_flipped_from_jagged_ijk."""
+        mock_jagged_ijk = MagicMock()
+        mock_morton_flipped.return_value = MagicMock()
+
+        space_filling_curve_from_jagged_ijk(mock_jagged_ijk, "z-trans")
+        mock_morton_flipped.assert_called_once_with(mock_jagged_ijk)
+
+    @patch("fvdb_extensions.models.fvdb_utils.morton_flipped_from_jagged_ijk")
+    def test_morton_flipped_alias_morton_zyx(self, mock_morton_flipped):
+        """Test that 'morton_zyx' dispatches to morton_flipped_from_jagged_ijk."""
+        mock_jagged_ijk = MagicMock()
+        mock_morton_flipped.return_value = MagicMock()
+
+        space_filling_curve_from_jagged_ijk(mock_jagged_ijk, "morton_zyx")
+        mock_morton_flipped.assert_called_once_with(mock_jagged_ijk)
+
+    @patch("fvdb_extensions.models.fvdb_utils.hilbert_from_jagged_ijk")
+    def test_hilbert_dispatch(self, mock_hilbert):
+        """Test that 'hilbert' dispatches to hilbert_from_jagged_ijk."""
+        mock_jagged_ijk = MagicMock()
+        mock_hilbert.return_value = MagicMock()
+
+        space_filling_curve_from_jagged_ijk(mock_jagged_ijk, "hilbert")
+        mock_hilbert.assert_called_once_with(mock_jagged_ijk)
+
+    @patch("fvdb_extensions.models.fvdb_utils.hilbert_flipped_from_jagged_ijk")
+    def test_hilbert_trans_dispatch(self, mock_hilbert_flipped):
+        """Test that 'hilbert-trans' dispatches to hilbert_flipped_from_jagged_ijk."""
+        mock_jagged_ijk = MagicMock()
+        mock_hilbert_flipped.return_value = MagicMock()
+
+        space_filling_curve_from_jagged_ijk(mock_jagged_ijk, "hilbert-trans")
+        mock_hilbert_flipped.assert_called_once_with(mock_jagged_ijk)
+
+    @patch("fvdb_extensions.models.fvdb_utils.identity_from_jagged_ijk")
+    def test_identity_alias_vdb(self, mock_identity):
+        """Test that 'vdb' dispatches to identity_from_jagged_ijk."""
+        mock_jagged_ijk = MagicMock()
+        mock_identity.return_value = MagicMock()
+
+        space_filling_curve_from_jagged_ijk(mock_jagged_ijk, "vdb")
+        mock_identity.assert_called_once_with(mock_jagged_ijk)
+
+    @patch("fvdb_extensions.models.fvdb_utils.identity_from_jagged_ijk")
+    def test_identity_alias_identity(self, mock_identity):
+        """Test that 'identity' dispatches to identity_from_jagged_ijk."""
+        mock_jagged_ijk = MagicMock()
+        mock_identity.return_value = MagicMock()
+
+        space_filling_curve_from_jagged_ijk(mock_jagged_ijk, "identity")
+        mock_identity.assert_called_once_with(mock_jagged_ijk)
+
+
+class TestValidCurveTypes(unittest.TestCase):
+    """Test that all documented curve types are valid."""
+
+    def test_all_valid_curve_types_documented(self):
+        """Verify the complete set of valid curve types."""
+        valid_types = {
+            "z",
+            "morton",
+            "z-trans",
+            "morton_zyx",
+            "hilbert",
+            "hilbert-trans",
+            "vdb",
+            "identity",
+        }
+        # This test documents the expected valid types
+        # The dispatch tests above verify each one works
+        self.assertEqual(len(valid_types), 8, "Should have 8 valid curve type aliases")
 
 
 if __name__ == "__main__":
     unittest.main()
-