diff --git a/tests/test_issue_40_cli_integration.py b/tests/test_issue_40_cli_integration.py
index 23674b6a..9992ae7c 100644
--- a/tests/test_issue_40_cli_integration.py
+++ b/tests/test_issue_40_cli_integration.py
@@ -96,9 +96,11 @@ class TestIssue40CLIIntegration:
         # Should create associated markdown file
         assert expected_md.exists()
 
-        # Verify it's valid markdown
+        # Verify it's valid markdown with schema reference
         md_content = expected_md.read_text()
-        assert md_content.startswith('# ')
+        # Content should include schema reference metadata and heading
+        assert '<!-- Generated from schema:' in md_content
+        assert '# Template Schema' in md_content
 
     def test_generate_stub_interactive_mode_defaults_to_associated_path(self, runner, temp_dir):
         """generate-stub in interactive mode should default to associated path."""
@@ -129,9 +131,11 @@ class TestIssue40CLIIntegration:
         expected_md = temp_dir / "template.md"
         assert expected_md.exists()
 
-        # Verify it's valid markdown
+        # Verify it's valid markdown with schema reference
         md_content = expected_md.read_text()
-        assert md_content.startswith('# ')
+        # Content should include schema reference metadata and heading
+        assert '<!-- Generated from schema:' in md_content
+        assert '# Template Schema' in md_content
 
     def test_validate_auto_discovers_associated_schema(self, runner, temp_dir):
         """validate should auto-discover associated schema when not specified."""
diff --git a/tests/test_issue_55_schema_based_draft_generation.py b/tests/test_issue_55_schema_based_draft_generation.py
index 37c71966..dd1d0ac0 100644
--- a/tests/test_issue_55_schema_based_draft_generation.py
+++ b/tests/test_issue_55_schema_based_draft_generation.py
@@ -382,7 +382,7 @@ class TestIssue55SchemaBasedDraftGeneration:
             validate_result = self.runner.invoke(cli, [
                 'validate',
                 str(draft_file),
-                str(schema_file)
+                '--schema', str(schema_file)
             ])
 
             # Assert - Generated draft should be valid against its source schema