Fix pytest tests to exit with proper exit code (#16681)

test/registered/layers/mamba/test_causal_conv1d.py

@@ -380,4 +380,6 @@ def test_causal_conv1d_varlen(
 
 
 if __name__ == "__main__":
-    pytest.main([__file__])
+    import sys
+
+    sys.exit(pytest.main([__file__]))

test/registered/layers/mamba/test_mamba2_mixer.py

@@ -144,4 +144,6 @@ def mixer2_gated_norm_tensor_parallel(
 
 
 if __name__ == "__main__":
-    pytest.main([__file__])
+    import sys
+
+    sys.exit(pytest.main([__file__]))

test/registered/layers/mamba/test_mamba_ssm.py

@@ -297,4 +297,6 @@ def test_selective_state_update_with_heads_with_batch_indices(
 
 
 if __name__ == "__main__":
-    pytest.main([__file__])
+    import sys
+
+    sys.exit(pytest.main([__file__]))

test/registered/layers/mamba/test_mamba_ssm_ssd.py

@@ -607,4 +607,6 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
 
 
 if __name__ == "__main__":
-    pytest.main([__file__])
+    import sys
+
+    sys.exit(pytest.main([__file__]))