MLIA-963: Capture and handle Vela warning

Signed-off-by: Gergely Nagy <gergely.nagy@arm.com>
Change-Id: I6c8b0b74d6d35261eb0ff1a37b9577f9033be8f9

1 files changed, 20 insertions, 0 deletions

diff --git a/tests/test_backend_vela_compiler.py b/tests/test_backend_vela_compiler.py
index 9b69ada..9f09efb 100644
--- a/tests/test_backend_vela_compiler.py
+++ b/tests/test_backend_vela_compiler.py
@@ -2,7 +2,9 @@
 # SPDX-License-Identifier: Apache-2.0
 """Tests for module vela/compiler."""
 from pathlib import Path
+from typing import Any
 
+import pytest
 from ethosu.vela.compiler_driver import TensorAllocator
 from ethosu.vela.scheduler import OptimizationStrategy
 
@@ -154,6 +156,24 @@ def test_compile_model(test_tflite_model: Path) -> None:
     assert isinstance(optimized_model, OptimizedModel)
 
 
+def test_compile_model_fail_sram_exceeded(
+    test_tflite_model: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Test model optimization."""
+    compiler = VelaCompiler(
+        EthosUConfiguration.load_profile("ethos-u55-256").compiler_options
+    )
+
+    def fake_compiler(*_: Any) -> None:
+        print("Warning: SRAM target for arena memory area exceeded.")
+
+    monkeypatch.setattr("mlia.backend.vela.compiler.compiler_driver", fake_compiler)
+    with pytest.raises(Exception) as exc_info:
+        compiler.compile_model(test_tflite_model)
+
+    assert str(exc_info.value) == "Model is too large and uses too much RAM"
+
+
 def test_optimize_model(tmp_path: Path, test_tflite_model: Path) -> None:
     """Test model optimization and saving into file."""
     tmp_file = tmp_path / "temp.tflite"