Skip to content

Commit 377b51c

Browse files
authored
test: microgen - adds tests to confirm attribute extraction from classes (#2315)
* test: adds tests to confirm attribute extraction from classes * Update scripts/microgenerator/tests/unit/test_generate_analyzer.py * adds docstring
1 parent a309a79 commit 377b51c

File tree

1 file changed

+172
-0
lines changed

1 file changed

+172
-0
lines changed

scripts/microgenerator/tests/unit/test_generate_analyzer.py

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,3 +87,175 @@ def test_import_extraction(self, code_snippet, expected_imports):
8787
expected = sorted(expected_imports)
8888

8989
assert extracted == expected
90+
91+
92+
# --- Tests CodeAnalyzer handling of Attributes ---
93+
94+
95+
class TestCodeAnalyzerAttributes:
96+
97+
@pytest.mark.parametrize(
98+
"code_snippet, expected_structure",
99+
[
100+
pytest.param(
101+
"""
102+
class MyClass:
103+
CLASS_VAR = 123
104+
""",
105+
[
106+
{
107+
"class_name": "MyClass",
108+
"methods": [],
109+
"attributes": [{"name": "CLASS_VAR", "type": None}],
110+
}
111+
],
112+
id="class_var_assign",
113+
),
114+
pytest.param(
115+
"""
116+
class MyClass:
117+
class_var: int = 456
118+
""",
119+
[
120+
{
121+
"class_name": "MyClass",
122+
"methods": [],
123+
"attributes": [{"name": "class_var", "type": "int"}],
124+
}
125+
],
126+
id="class_var_annassign",
127+
),
128+
pytest.param(
129+
"""
130+
class MyClass:
131+
class_var: int
132+
""",
133+
[
134+
{
135+
"class_name": "MyClass",
136+
"methods": [],
137+
"attributes": [{"name": "class_var", "type": "int"}],
138+
}
139+
],
140+
id="class_var_annassign_no_value",
141+
),
142+
pytest.param(
143+
"""
144+
class MyClass:
145+
def __init__(self):
146+
self.instance_var = 789
147+
""",
148+
[
149+
{
150+
"class_name": "MyClass",
151+
"methods": [
152+
{
153+
"method_name": "__init__",
154+
"args": [{"name": "self", "type": None}],
155+
"return_type": None,
156+
}
157+
],
158+
"attributes": [{"name": "instance_var", "type": None}],
159+
}
160+
],
161+
id="instance_var_assign",
162+
),
163+
pytest.param(
164+
"""
165+
class MyClass:
166+
def __init__(self):
167+
self.instance_var: str = 'hello'
168+
""",
169+
[
170+
{
171+
"class_name": "MyClass",
172+
"methods": [
173+
{
174+
"method_name": "__init__",
175+
"args": [{"name": "self", "type": None}],
176+
"return_type": None,
177+
}
178+
],
179+
"attributes": [{"name": "instance_var", "type": "str"}],
180+
}
181+
],
182+
id="instance_var_annassign",
183+
),
184+
pytest.param(
185+
"""
186+
class MyClass:
187+
def __init__(self):
188+
self.instance_var: str
189+
""",
190+
[
191+
{
192+
"class_name": "MyClass",
193+
"methods": [
194+
{
195+
"method_name": "__init__",
196+
"args": [{"name": "self", "type": None}],
197+
"return_type": None,
198+
}
199+
],
200+
"attributes": [{"name": "instance_var", "type": "str"}],
201+
}
202+
],
203+
id="instance_var_annassign_no_value",
204+
),
205+
pytest.param(
206+
"""
207+
class MyClass:
208+
VAR_A = 1
209+
var_b: int = 2
210+
def __init__(self):
211+
self.var_c = 3
212+
self.var_d: float = 4.0
213+
""",
214+
[
215+
{
216+
"class_name": "MyClass",
217+
"methods": [
218+
{
219+
"method_name": "__init__",
220+
"args": [{"name": "self", "type": None}],
221+
"return_type": None,
222+
}
223+
],
224+
"attributes": [
225+
{"name": "VAR_A", "type": None},
226+
{"name": "var_b", "type": "int"},
227+
{"name": "var_c", "type": None},
228+
{"name": "var_d", "type": "float"},
229+
],
230+
}
231+
],
232+
id="mixed_attributes",
233+
),
234+
pytest.param(
235+
"a = 123 # Module level",
236+
[],
237+
id="module_level_assign",
238+
),
239+
pytest.param(
240+
"b: int = 456 # Module level",
241+
[],
242+
id="module_level_annassign",
243+
),
244+
],
245+
)
246+
def test_attribute_extraction(self, code_snippet: str, expected_structure: list):
247+
"""Tests the extraction of class and instance attributes."""
248+
analyzer = CodeAnalyzer()
249+
tree = ast.parse(code_snippet)
250+
analyzer.visit(tree)
251+
252+
extracted = analyzer.structure
253+
# Normalize attributes for order-independent comparison
254+
for item in extracted:
255+
if "attributes" in item:
256+
item["attributes"].sort(key=lambda x: x["name"])
257+
for item in expected_structure:
258+
if "attributes" in item:
259+
item["attributes"].sort(key=lambda x: x["name"])
260+
261+
assert extracted == expected_structure

0 commit comments

Comments
 (0)