refactor: change to check expected and not written file size

AllenNeuralDynamics · Jan 3, 2025 · 3dfd475 · 3dfd475
1 parent a2fe7ea
commit 3dfd475
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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