diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 7b0029c73cb..3d57691cea4 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -197,12 +197,18 @@ class Markers: if start <= marker.start < end: return True if marker.standalone: - return self.only_separators(marker.end, start) + return self.on_the_line_below(marker.end, start) return self.only_separators(end, marker.start) + def on_the_line_below(self, start: int, end: int) -> bool: + """Whether a marker on its own line is written directly above the statement, which means + one line break and nothing else that carries meaning. A blank line between the two leaves + the marker reading as a note about the file rather than a bound on what follows it.""" + return self.only_separators(start, end) and self.sql[start:end].count("\n") == 1 + def only_separators(self, start: int, end: int) -> bool: """Whether nothing but statement separators lie between two points, which is what makes a - marker and a statement adjacent whatever whitespace and line breaks sit between them.""" + marker and the statement it follows adjacent however they are laid out.""" return start <= end and not self.sql[start:end].strip(" \t\r\n;") diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index 1573a1b5706..5413f8cdc62 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -375,6 +375,11 @@ class TestEscapeHatch: assert _keywords(tmp_path, sql) == ("UPDATE",) assert _scan(tmp_path, sql)[0].line == 1 + def test_a_marker_a_blank_line_above_a_statement_does_not_exempt_it(self, tmp_path): + sql = '-- data-migration-ok: one row\n\nUPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + def test_a_trailing_marker_exempts_only_the_statement_it_follows(self, tmp_path): sql = 'DELETE FROM "Foo" WHERE "a" = 1; UPDATE "Bar" SET "b" = 2; -- data-migration-ok: one row' assert _keywords(tmp_path, sql) == ("DELETE",) @@ -410,6 +415,11 @@ class TestEscapeHatch: assert _keywords(tmp_path, sql) == ("UPDATE",) assert _scan(tmp_path, sql)[0].line == 5 + def test_marker_directly_above_a_one_line_do_block_does_not_exempt_its_body(self, tmp_path): + sql = '-- data-migration-ok: seeding one default row\nDO $$ BEGIN UPDATE "Foo" SET "a" = 1; END $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 2 + def test_marker_on_the_do_line_does_not_exempt_its_body(self, tmp_path): sql = ( "DO $$ -- data-migration-ok: bounded to one row\n"