final unit tests for hdf5

Author: arthurmloureiro

src/smokescreen/utils.py

@@ -157,7 +157,7 @@ def load_sacc_file(path_to_sacc: str) -> tuple[sacc.Sacc, str]:
     Returns
     -------
     sacc.Sacc
-        Loaded SACC object
+        Loaded SACC object with _smokescreen_input_format attribute set
     str
         Detected input format ('fits' or 'hdf5')
 
@@ -169,13 +169,15 @@ def load_sacc_file(path_to_sacc: str) -> tuple[sacc.Sacc, str]:
     # Try HDF5 first (more specific format check)
     try:
         sacc_obj = sacc.Sacc.load_hdf5(path_to_sacc)
+        sacc_obj._smokescreen_input_format = 'hdf5'
         return sacc_obj, 'hdf5'
     except Exception:
         pass
 
     # Fall back to FITS format
     try:
         sacc_obj = sacc.Sacc.load_fits(path_to_sacc)
+        sacc_obj._smokescreen_input_format = 'fits'
         return sacc_obj, 'fits'
     except Exception as e:
         raise ValueError(

tests/test_encryption.py

@@ -1,5 +1,6 @@
 import os
 import pytest
+from cryptography.fernet import Fernet
 from smokescreen.encryption import encrypt_file, decrypt_file
 
 
@@ -120,3 +121,64 @@ def test_decrypt_file_nonexistent_key(encrypted_file_and_key):
     # Test decrypting with a nonexistent key
     with pytest.raises(FileNotFoundError):
         decrypt_file(str(encrypted_file_path), "nonexistent_key.key")
+
+
+def test_decrypt_file_save_with_fallback_handling(tmp_path):
+    # Test the fallback handling in decrypt_file when filename doesn't end with .encrpt
+    # Note: encrypt_file stores as basename.encrpt where basename is split('.')[0]
+    # So data.sacc becomes data.encrpt, and decrypts back to data (original extension lost)
+    original_content = b"test content for decryption fallback"
+    original_file = tmp_path / "data.sacc"
+    original_file.write_bytes(original_content)
+
+    # Encrypt it - basename is extracted as 'data' from 'data.sacc'
+    encrypted_sacc, key = encrypt_file(str(original_file), path_to_save=str(tmp_path),
+                                       save_file=True, keep_original=False)
+    encrypted_file = tmp_path / "data.encrpt"
+    key_file = tmp_path / "data.key"
+
+    assert encrypted_file.exists()
+    assert key_file.exists()
+
+    # Decrypt with save_file=True
+    decrypt_file(str(encrypted_file), str(key_file), save_file=True)
+
+    # The decrypted file restores basename (without .encrpt) which is 'data'
+    # Note: original extension '.sacc' cannot be recovered after encryption
+    decrypted_file = tmp_path / "data"
+    assert decrypted_file.exists()
+
+    # Verify content matches original
+    assert decrypted_file.read_bytes() == original_content
+
+
+def test_decrypt_file_save_with_fallback_extension_removal(tmp_path):
+    # Test fallback handling when .encrpt appears in middle of filename
+    # Note: encrypt_file stores as basename.encrpt where basename is split('.')[0]
+    # So data.backup becomes data.encrpt, and decrypts back to data (original extension lost)
+
+    original_content = b"test content for edge case decryption"
+    backup_content = b"backup content that was encrypted"
+
+    # Encrypt a backup file
+    backup_file = tmp_path / "data.backup"
+    backup_file.write_bytes(backup_content)
+
+    encrypted_sacc, key = encrypt_file(str(backup_file), path_to_save=str(tmp_path),
+                                       save_file=True, keep_original=False)
+    # The encrypted file is data.encrpt (basename extracted from 'data.backup')
+    assert (tmp_path / "data.encrpt").exists()
+
+    # Get the key file path
+    key_file = tmp_path / "data.key"
+
+    # Decrypt it using the key file
+    decrypted_sacc = decrypt_file(str(tmp_path / "data.encrpt"), str(key_file), save_file=True)
+
+    # Should restore to basename (without .encrpt) which is 'data'
+    # Note: original extension '.backup' cannot be recovered after encryption
+    decrypted_file = tmp_path / "data"
+    assert decrypted_file.exists()
+    assert decrypted_file.read_bytes() == backup_content
+
+

tests/test_utils.py

@@ -77,8 +77,9 @@ class TestLoadSaccFile:
     """Tests for the load_sacc_file utility function."""
 
     def test_load_sacc_file_fits_format(self, tmp_path):
+        import sacc as sacc_mod
         # Create a FITS SACC file
-        sacc_data = sacc.Sacc()
+        sacc_data = sacc_mod.Sacc()
         sacc_data.add_tracer('misc', 'test')
         sacc_data.add_data_point('galaxy_shear_cl_ee', ('test', 'test'), 1.0, ell=10)
         fits_path = tmp_path / "test.fits"
@@ -87,13 +88,14 @@ def test_load_sacc_file_fits_format(self, tmp_path):
         # Load using load_sacc_file
         loaded_sacc, file_format = load_sacc_file(str(fits_path))
 
-        assert isinstance(loaded_sacc, sacc.Sacc)
+        assert isinstance(loaded_sacc, sacc_mod.Sacc)
         assert file_format == 'fits'
         assert len(loaded_sacc.mean) == 1
 
     def test_load_sacc_file_hdf5_format(self, tmp_path):
+        import sacc as sacc_mod
         # Create an HDF5 SACC file
-        sacc_data = sacc.Sacc()
+        sacc_data = sacc_mod.Sacc()
         sacc_data.add_tracer('misc', 'test')
         sacc_data.add_data_point('galaxy_shear_cl_ee', ('test', 'test'), 1.0, ell=10)
         hdf5_path = tmp_path / "test.hdf5"
@@ -102,13 +104,14 @@ def test_load_sacc_file_hdf5_format(self, tmp_path):
         # Load using load_sacc_file
         loaded_sacc, file_format = load_sacc_file(str(hdf5_path))
 
-        assert isinstance(loaded_sacc, sacc.Sacc)
+        assert isinstance(loaded_sacc, sacc_mod.Sacc)
         assert file_format == 'hdf5'
         assert len(loaded_sacc.mean) == 1
 
     def test_load_sacc_file_with_h5_extension(self, tmp_path):
+        import sacc as sacc_mod
         # Create an HDF5 SACC file with .h5 extension
-        sacc_data = sacc.Sacc()
+        sacc_data = sacc_mod.Sacc()
         sacc_data.add_tracer('misc', 'test')
         sacc_data.add_data_point('galaxy_shear_cl_ee', ('test', 'test'), 1.0, ell=10)
         h5_path = tmp_path / "test.h5"
@@ -117,12 +120,13 @@ def test_load_sacc_file_with_h5_extension(self, tmp_path):
         # Load using load_sacc_file - should detect as HDF5 regardless of extension
         loaded_sacc, file_format = load_sacc_file(str(h5_path))
 
-        assert isinstance(loaded_sacc, sacc.Sacc)
+        assert isinstance(loaded_sacc, sacc_mod.Sacc)
         assert file_format == 'hdf5'
 
     def test_load_sacc_file_with_sacc_extension_hdf5(self, tmp_path):
+        import sacc as sacc_mod
         # Create an HDF5 SACC file with .sacc extension (like sn_datavector.sacc)
-        sacc_data = sacc.Sacc()
+        sacc_data = sacc_mod.Sacc()
         sacc_data.add_tracer('misc', 'test')
         sacc_data.add_data_point('galaxy_shear_cl_ee', ('test', 'test'), 1.0, ell=10)
         sacc_path = tmp_path / "test.sacc"
@@ -131,7 +135,7 @@ def test_load_sacc_file_with_sacc_extension_hdf5(self, tmp_path):
         # Load using load_sacc_file - should detect as HDF5 even with .sacc extension
         loaded_sacc, file_format = load_sacc_file(str(sacc_path))
 
-        assert isinstance(loaded_sacc, sacc.Sacc)
+        assert isinstance(loaded_sacc, sacc_mod.Sacc)
         assert file_format == 'hdf5'
 
     def test_load_sacc_file_nonexistent(self):