1+ from __future__ import annotations
2+
3+ import os
4+ import inspect
5+ import typing as t
6+ from unittest .mock import MagicMock
7+
8+ import pytest
9+ from sqlmesh .core .model import Model
10+ from sqlmesh .core .linter .rule import Rule , RuleViolation
11+
12+
13+ class TestRule (Rule ):
14+ """A test rule for testing the get_definition_location method."""
15+
16+ def check_model (self , model : Model ) -> t .Optional [RuleViolation ]:
17+ """The evaluation function that'll check for a violation of this rule."""
18+ return None
19+
20+
21+ def test_get_definition_location ():
22+ """Test the get_definition_location method returns correct file and line information."""
23+ # Create a mock context
24+ mock_context = MagicMock ()
25+ rule = TestRule (mock_context )
26+
27+ # Get the expected location using the inspect module
28+ expected_file = inspect .getfile (TestRule )
29+ expected_source , expected_start_line = inspect .getsourcelines (TestRule )
30+ expected_end_line = expected_start_line + len (expected_source ) - 1
31+
32+ # Get the location using the Rule method
33+ location = rule .get_definition_location ()
34+
35+ # Assert the file path matches
36+ assert location ["file_path" ] == expected_file
37+
38+ # Assert the line numbers match
39+ assert location ["start_line" ] == expected_start_line
40+ assert location ["end_line" ] == expected_end_line
41+
42+ # Test the fallback case for a class without source
43+ with pytest .MonkeyPatch .context () as mp :
44+ # Mock inspect.getsourcelines to raise an exception
45+ def mock_getsourcelines (* args , ** kwargs ):
46+ raise IOError ("Mock error" )
47+
48+ mp .setattr (inspect , "getsourcelines" , mock_getsourcelines )
49+
50+ # Get the location with the mocked function
51+ fallback_location = rule .get_definition_location ()
52+
53+ # It should still have the file path
54+ assert "file_path" in fallback_location
55+ assert fallback_location ["file_path" ] == expected_file
56+
57+ # But not the line numbers
58+ assert "start_line" not in fallback_location
59+ assert "end_line" not in fallback_location
0 commit comments