diff --git a/oocana/oocana/serialization.py b/oocana/oocana/serialization.py index cedd49c..ec43093 100644 --- a/oocana/oocana/serialization.py +++ b/oocana/oocana/serialization.py @@ -69,28 +69,28 @@ def compression_options(context: 'Context') -> CompressionOptions | None: print(f"An unexpected error occurred while reading compression options: {e}. Returning None.") return None +# Mapping from compression method to file suffix +COMPRESSION_SUFFIXES = { + "zip": ".zip", + "gzip": ".gz", + "bz2": ".bz2", + "zstd": ".zst", + "xz": ".xz", + "tar": ".tar", +} + def compression_suffix(context: 'Context') -> str: """ Get the file suffix based on the compression method. - If no compression is specified, return an empty string. + If no compression is specified, return ".pkl" (pickle format). """ compression = compression_options(context) - if compression is None or compression["method"] is None: + if compression is None: return ".pkl" - - method = compression["method"] - if method == "zip": - return ".zip" - elif method == "gzip": - return ".gz" - elif method == "bz2": - return ".bz2" - elif method == "zstd": - return ".zst" - elif method == "xz": - return ".xz" - elif method == "tar": - return ".tar" - else: - return ".pkl" # Default case if method is not recognized \ No newline at end of file + + method = compression.get("method") + if method is None: + return ".pkl" + + return COMPRESSION_SUFFIXES.get(method, ".pkl") \ No newline at end of file diff --git a/oocana/tests/test_serialization.py b/oocana/tests/test_serialization.py new file mode 100644 index 0000000..5d5afa3 --- /dev/null +++ b/oocana/tests/test_serialization.py @@ -0,0 +1,164 @@ +import unittest +from unittest.mock import MagicMock, patch +import tempfile +import os +import json + + +class TestCompressionSuffix(unittest.TestCase): + """Test cases for compression_suffix function.""" + + def setUp(self): + # Create a temporary directory for the mock context + self.temp_dir = tempfile.mkdtemp() + self.mock_context = MagicMock() + self.mock_context.pkg_data_dir = self.temp_dir + + def tearDown(self): + # Clean up temporary files + import shutil + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_compression_suffix_no_options(self): + """Test that .pkl is returned when no compression options exist.""" + from oocana.oocana.serialization import compression_suffix + + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".pkl") + + def test_compression_suffix_zip(self): + """Test compression suffix for zip method.""" + from oocana.oocana.serialization import compression_suffix, COMPRESSION_OPTIONS_FILE + + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump({"method": "zip"}, f) + + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".zip") + + def test_compression_suffix_gzip(self): + """Test compression suffix for gzip method.""" + from oocana.oocana.serialization import compression_suffix, COMPRESSION_OPTIONS_FILE + + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump({"method": "gzip"}, f) + + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".gz") + + def test_compression_suffix_bz2(self): + """Test compression suffix for bz2 method.""" + from oocana.oocana.serialization import compression_suffix, COMPRESSION_OPTIONS_FILE + + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump({"method": "bz2"}, f) + + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".bz2") + + def test_compression_suffix_zstd(self): + """Test compression suffix for zstd method.""" + from oocana.oocana.serialization import compression_suffix, COMPRESSION_OPTIONS_FILE + + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump({"method": "zstd"}, f) + + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".zst") + + def test_compression_suffix_xz(self): + """Test compression suffix for xz method.""" + from oocana.oocana.serialization import compression_suffix, COMPRESSION_OPTIONS_FILE + + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump({"method": "xz"}, f) + + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".xz") + + def test_compression_suffix_tar(self): + """Test compression suffix for tar method.""" + from oocana.oocana.serialization import compression_suffix, COMPRESSION_OPTIONS_FILE + + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump({"method": "tar"}, f) + + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".tar") + + def test_compression_suffix_null_method(self): + """Test that .pkl is returned when method is null.""" + from oocana.oocana.serialization import compression_suffix, COMPRESSION_OPTIONS_FILE + + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump({"method": None}, f) + + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".pkl") + + def test_compression_suffix_missing_method_key(self): + """Test that .pkl is returned when method key is missing (no KeyError).""" + from oocana.oocana.serialization import compression_suffix, COMPRESSION_OPTIONS_FILE + + # Write a dict without the 'method' key + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump({"other_key": "value"}, f) + + # Should not raise KeyError + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".pkl") + + def test_compression_suffix_unknown_method(self): + """Test that .pkl is returned for unknown compression method.""" + from oocana.oocana.serialization import compression_suffix, COMPRESSION_OPTIONS_FILE + + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump({"method": "unknown_method"}, f) + + result = compression_suffix(self.mock_context) + self.assertEqual(result, ".pkl") + + +class TestCompressionOptions(unittest.TestCase): + """Test cases for compression_options function.""" + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.mock_context = MagicMock() + self.mock_context.pkg_data_dir = self.temp_dir + + def tearDown(self): + import shutil + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_compression_options_returns_none_when_no_file(self): + """Test that None is returned when options file doesn't exist.""" + from oocana.oocana.serialization import compression_options + + result = compression_options(self.mock_context) + self.assertIsNone(result) + + def test_compression_options_returns_dict_when_file_exists(self): + """Test that options are returned when file exists.""" + from oocana.oocana.serialization import compression_options, COMPRESSION_OPTIONS_FILE + + expected = {"method": "gzip"} + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + json.dump(expected, f) + + result = compression_options(self.mock_context) + self.assertEqual(result, expected) + + def test_compression_options_handles_invalid_json(self): + """Test that None is returned for invalid JSON.""" + from oocana.oocana.serialization import compression_options, COMPRESSION_OPTIONS_FILE + + with open(os.path.join(self.temp_dir, COMPRESSION_OPTIONS_FILE), "w") as f: + f.write("not valid json") + + result = compression_options(self.mock_context) + self.assertIsNone(result) + + +if __name__ == '__main__': + unittest.main()