Fix resource parsing

src/lib/utils/tests/BUILD

@@ -48,5 +48,14 @@ osmo_py_test(
     srcs = ["test_redact_secrets.py"],
     deps = [
         "//src/lib/utils:redact",
-    ],
+    ]
+)
+
+osmo_py_test(
+    name = "test_workflow",
+    srcs = ["test_workflow.py"],
+    deps = [
+        "//src/lib/utils:osmo_errors",
+        "//src/lib/utils:workflow",
+    ]
 )

src/lib/utils/tests/test_workflow.py

@@ -0,0 +1,109 @@
+"""
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+SPDX-License-Identifier: Apache-2.0
+"""
+import unittest
+
+from src.lib.utils import osmo_errors, workflow as workflow_utils
+
+
+class ParseWorkflowSpecTests(unittest.TestCase):
+    """Unit tests for parse_workflow_spec."""
+
+    def test_workflow_only(self):
+        spec = """\
+workflow:
+  name: my-wf
+  groups: []
+"""
+        workflow_spec, default_values = workflow_utils.parse_workflow_spec(spec)
+        self.assertIn('workflow:', workflow_spec)
+        self.assertIsNone(default_values)
+
+    def test_workflow_then_default_values(self):
+        spec = """\
+workflow:
+  name: my-wf
+  groups: []
+default-values:
+  foo: bar
+"""
+        workflow_spec, default_values = workflow_utils.parse_workflow_spec(spec)
+        self.assertIn('name: my-wf', workflow_spec)
+        self.assertEqual(default_values, {'foo': 'bar'})
+
+    def test_default_values_before_workflow(self):
+        spec = """\
+default-values:
+  foo: bar
+workflow:
+  name: my-wf
+  groups: []
+"""
+        workflow_spec, default_values = workflow_utils.parse_workflow_spec(spec)
+        self.assertIn('name: my-wf', workflow_spec)
+        self.assertEqual(default_values, {'foo': 'bar'})
+
+    def test_jinja_content_not_at_root_indent(self):
+        spec = """\
+workflow:
+  name: my-wf
+  groups:
+{% for i in range(3) %}
+  - name: task-{{ i }}
+{% endfor %}
+"""
+        workflow_spec, default_values = workflow_utils.parse_workflow_spec(spec)
+        self.assertIn('task-', workflow_spec)
+        self.assertIsNone(default_values)
+
+    def test_duplicate_workflow_raises(self):
+        spec = """\
+workflow:
+  name: first
+workflow:
+  name: second
+"""
+        with self.assertRaises(osmo_errors.OSMOUserError) as context:
+            workflow_utils.parse_workflow_spec(spec)
+        self.assertIn('workflow', str(context.exception))
+
+    def test_version_key_allowed(self):
+        spec = """\
+version: 2
+workflow:
+  name: my-wf
+  groups: []
+"""
+        workflow_spec, default_values = workflow_utils.parse_workflow_spec(spec)
+        self.assertIn('name: my-wf', workflow_spec)
+        self.assertIsNone(default_values)
+
+    def test_unknown_top_level_key_raises(self):
+        spec = """\
+workflow:
+  name: my-wf
+resources:
+  default:
+    cpu: 10
+"""
+        with self.assertRaises(osmo_errors.OSMOUserError) as context:
+            workflow_utils.parse_workflow_spec(spec)
+        self.assertIn('resources', str(context.exception))
+
+
+if __name__ == '__main__':
+    unittest.main()

src/lib/utils/workflow.py

@@ -41,28 +41,27 @@ def fetch_default_values(workflow_spec: str) -> str | None:
 
 def parse_workflow_spec(workflow_spec: str) -> Tuple[str, Dict | None]:
     """ Parse the workflow spec. """
-    workflow_pattern = re.compile(r'(^workflow:.*?)(?=^workflow:|\Z)',
-                                  re.DOTALL | re.MULTILINE)
-    workflows = workflow_pattern.findall(workflow_spec)
-    if len(workflows) > 1:
-        raise osmo_errors.OSMOUserError('Multiple workflows sections found in the workflow spec.')
-
-    if workflows:
-        workflow_portion = workflows[0]
-    else:
+    section_pattern = re.compile(
+        r'^([a-zA-Z][a-zA-Z0-9_-]*):(.*?)(?=^[a-zA-Z][a-zA-Z0-9_-]*:|\Z)',
+        re.DOTALL | re.MULTILINE,
+    )
+    allowed_sections = {'workflow', 'default-values', 'version'}
+    sections: Dict[str, str] = {}
+    for match in section_pattern.finditer(workflow_spec):
+        key = match.group(1)
+        if key not in allowed_sections:
+            raise osmo_errors.OSMOUserError(
+                f'Unknown top-level key "{key}" found in the workflow spec.')
+        if key in sections:
+            raise osmo_errors.OSMOUserError(
+                f'Duplicate top-level key "{key}" found in the workflow spec.')
+        sections[key] = match.group(1) + ':' + match.group(2)
+
+    if 'workflow' not in sections:
         raise osmo_errors.OSMOUserError('Workflow spec not found.')
 
-    default_values_pattern = re.compile(r'(^default-values:.*?)(?=^(?![#\s])\S|\Z)',
-                                        re.DOTALL | re.MULTILINE)
-    default_locs = default_values_pattern.findall(workflow_spec)
-    if len(default_locs) > 1:
-        raise osmo_errors.OSMOUserError(
-            'Multiple default-values sections found in the workflow spec.')
-
     default_values = None
+    if 'default-values' in sections:
+        default_values = yaml.safe_load(sections['default-values'])['default-values']
 
-    # Get default values from 'default-values'
-    if default_locs:
-        default_values = yaml.safe_load(default_locs[0])['default-values']
-
-    return workflow_portion, default_values
+    return sections['workflow'], default_values

src/service/core/workflow/objects.py

@@ -813,8 +813,6 @@ def construct_workflow_dict(self, template_spec: workflow.TemplateSpec) -> Dict:
         try:
             updated_workflow_txt = template_spec.load_template_with_variables()
             updated_workflow_dict: Dict[str, Any] = yaml.safe_load(updated_workflow_txt)
-            if 'default-values' in updated_workflow_dict:
-                del updated_workflow_dict['default-values']
         except yaml.YAMLError as yaml_error:
             err_msg=f'Workflow spec is not properly formatted: {yaml_error}'
             # Construct a workflow ID with format <name>-<number>