From 3dfd475ac3a43ab7a28a94ed3bb19535af838ba0 Mon Sep 17 00:00:00 2001 From: Dan Birman Date: Fri, 3 Jan 2025 09:54:10 -0800 Subject: [PATCH] refactor: change to check expected and not written file size --- src/aind_data_schema/base.py | 3 ++- tests/test_base.py | 8 ++------ tests/test_examples.py | 4 +--- 3 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/aind_data_schema/base.py b/src/aind_data_schema/base.py index 806bc7d6c..6e0a61210 100644 --- a/src/aind_data_schema/base.py +++ b/src/aind_data_schema/base.py @@ -184,5 +184,6 @@ def write_standard_file( with open(filename, "w") as f: f.write(self.model_dump_json(indent=3)) - if os.path.getsize(filename) > MAX_FILE_SIZE: + # Check that size doesn't exceed the maximum + if len(self.model_dump_json(indent=3)) > MAX_FILE_SIZE: logging.warning(f"File size exceeds {MAX_FILE_SIZE / 1024} KB: {filename}") diff --git a/tests/test_base.py b/tests/test_base.py index d63e3d238..215d897f3 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -159,17 +159,13 @@ class Modelv2(AindCoreModel): # this is to ensure you can't get a bumped schema_version without passing validation self.assertRaises(ValidationError, lambda: Modelv1(**v2_from_v1.model_dump())) - @patch("os.path.getsize") @patch("builtins.open", new_callable=mock_open) @patch("logging.warning") - def test_write_standard_file_size_warning( - self, mock_logging_warning: MagicMock, mock_open: MagicMock, mock_getsize: MagicMock - ): + def test_write_standard_file_size_warning(self, mock_logging_warning: MagicMock, mock_open: MagicMock): """Tests that a warning is logged if the file size exceeds MAX_FILE_SIZE""" - mock_getsize.return_value = MAX_FILE_SIZE + 1 # Simulate file size exceeding the limit - s = Subject.model_construct() + s.subject_id = "s"*(MAX_FILE_SIZE + 1000) s.write_standard_file(output_directory=Path("dir"), suffix=".foo.bar") mock_open.assert_has_calls([call(Path("dir/subject.foo.bar"), "w")]) diff --git a/tests/test_examples.py b/tests/test_examples.py index 39dccbbf2..8e8e5d4f2 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -32,9 +32,7 @@ def test_examples(self): module = importlib.util.module_from_spec(spec) sys.modules["test_module"] = module - with patch("os.path.getsize", return_value=0), patch( - "builtins.open", new_callable=mock_open - ) as mocked_file: + with patch("builtins.open", new_callable=mock_open) as mocked_file: spec.loader.exec_module(module) h = mocked_file.return_value.__enter__() call_args_list = h.write.call_args_list