Skip to content

Commit e5ab1a7

Browse files
authored
[hec-assembler]: New unit tests for he_as.py and kernel splitter.py (#153)
* New unit tests for he_as and kernel splitter
1 parent 5d6ddcc commit e5ab1a7

File tree

9 files changed

+973
-69
lines changed

9 files changed

+973
-69
lines changed

assembler_tools/hec-assembler-tools/tests/unit_tests/test_assembler/test_prep/test_kernel_splitter.py

Lines changed: 476 additions & 0 deletions
Large diffs are not rendered by default.
Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
# Copyright (C) 2025 Intel Corporation
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
# These contents may have been developed with support from one or more Intel-operated
5+
# generative artificial intelligence solutions
6+
7+
import os
8+
import sys
9+
from unittest import mock
10+
11+
import he_as
12+
import pytest
13+
14+
15+
def test_parse_args_parses_all_flags(monkeypatch):
16+
test_args = [
17+
"prog",
18+
"kernel.tw",
19+
"--isa_spec",
20+
"isa.json",
21+
"--mem_spec",
22+
"mem.json",
23+
"--input_mem_file",
24+
"custom.mem",
25+
"--output_dir",
26+
"out",
27+
"--output_prefix",
28+
"pref",
29+
"--spad_size",
30+
"32",
31+
"--hbm_size",
32+
"64",
33+
"--no_hbm",
34+
"--repl_policy",
35+
he_as.constants.Constants.REPLACEMENT_POLICIES[0],
36+
"--use_xinstfetch",
37+
"--suppress_comments",
38+
"-vv",
39+
]
40+
monkeypatch.setattr(sys, "argv", test_args)
41+
args = he_as.parse_args()
42+
43+
assert args.input_file == "kernel.tw"
44+
assert args.isa_spec_file == "isa.json"
45+
assert args.mem_spec_file == "mem.json"
46+
assert args.input_mem_file == "custom.mem"
47+
assert args.output_dir == "out"
48+
assert args.output_prefix == "pref"
49+
assert args.spad_size == 32
50+
assert args.hbm_size == 64
51+
assert args.has_hbm is False
52+
assert args.repl_policy == he_as.constants.Constants.REPLACEMENT_POLICIES[0]
53+
assert args.use_xinstfetch is True
54+
assert args.suppress_comments is True
55+
assert args.debug_verbose == 2
56+
57+
58+
def test_run_config_derives_mem_file(tmp_path):
59+
input_file = tmp_path / "kernel.tw"
60+
input_file.write_text("")
61+
config = he_as.AssemblerRunConfig(input_file=str(input_file))
62+
63+
assert config.input_file == str(input_file)
64+
assert config.input_mem_file == str(input_file.with_suffix(".mem"))
65+
assert config.output_dir == os.path.dirname(os.path.realpath(str(input_file)))
66+
assert config.input_prefix == "kernel"
67+
68+
69+
def test_main_invokes_assembler_and_creates_outputs(tmp_path, monkeypatch):
70+
input_file = tmp_path / "kernel.tw"
71+
input_file.write_text("dummy")
72+
mem_file = tmp_path / "kernel.mem"
73+
mem_file.write_text("mem")
74+
output_dir = tmp_path / "outputs"
75+
config = he_as.AssemblerRunConfig(
76+
input_file=str(input_file),
77+
input_mem_file=str(mem_file),
78+
output_dir=str(output_dir),
79+
output_prefix="result",
80+
)
81+
82+
asm_mock = mock.Mock(return_value=(1, 2, 3, 4, 5))
83+
monkeypatch.setattr(he_as, "asmisaAssemble", asm_mock)
84+
he_as.main(config, verbose=False)
85+
86+
assert asm_mock.called
87+
copied_config = asm_mock.call_args.args[0]
88+
assert copied_config is not config
89+
for ext in ("minst", "cinst", "xinst"):
90+
assert (tmp_path / f"outputs/result.{ext}").is_file()
91+
92+
93+
def test_run_config_requires_input_file():
94+
"""Test that AssemblerRunConfig raises TypeError when input_file is missing."""
95+
with pytest.raises(TypeError, match="Expected value for configuration `input_file`"):
96+
he_as.AssemblerRunConfig()
97+
98+
99+
def test_run_config_defaults(tmp_path):
100+
"""Test that AssemblerRunConfig sets sensible defaults."""
101+
input_file = tmp_path / "kernel.tw"
102+
input_file.write_text("")
103+
config = he_as.AssemblerRunConfig(input_file=str(input_file))
104+
105+
assert config.has_hbm is True
106+
assert config.hbm_size == he_as.AssemblerRunConfig.DEFAULT_HBM_SIZE_KB
107+
assert config.spad_size == he_as.AssemblerRunConfig.DEFAULT_SPAD_SIZE_KB
108+
assert config.repl_policy == he_as.AssemblerRunConfig.DEFAULT_REPL_POLICY
109+
110+
111+
def test_run_config_custom_output_dir(tmp_path):
112+
"""Test that custom output_dir is respected."""
113+
input_file = tmp_path / "kernel.tw"
114+
input_file.write_text("")
115+
custom_dir = tmp_path / "custom_output"
116+
117+
config = he_as.AssemblerRunConfig(
118+
input_file=str(input_file),
119+
output_dir=str(custom_dir),
120+
)
121+
122+
assert config.output_dir == str(custom_dir)
123+
124+
125+
def test_main_creates_output_directory(tmp_path, monkeypatch):
126+
"""Test that main creates output directory if it doesn't exist."""
127+
input_file = tmp_path / "kernel.tw"
128+
input_file.write_text("dummy")
129+
mem_file = tmp_path / "kernel.mem"
130+
mem_file.write_text("mem")
131+
output_dir = tmp_path / "new_outputs"
132+
133+
config = he_as.AssemblerRunConfig(
134+
input_file=str(input_file),
135+
input_mem_file=str(mem_file),
136+
output_dir=str(output_dir),
137+
)
138+
139+
asm_mock = mock.Mock(return_value=(1, 2, 3, 4, 5))
140+
monkeypatch.setattr(he_as, "asmisaAssemble", asm_mock)
141+
142+
assert not output_dir.exists()
143+
he_as.main(config, verbose=False)
144+
assert output_dir.exists()
145+
146+
147+
def test_main_output_not_writable_raises_exception(tmp_path, monkeypatch):
148+
"""Test that main raises exception when output location is not writable."""
149+
input_file = tmp_path / "kernel.tw"
150+
input_file.write_text("dummy")
151+
mem_file = tmp_path / "kernel.mem"
152+
mem_file.write_text("mem")
153+
output_dir = tmp_path / "outputs"
154+
output_dir.mkdir()
155+
os.chmod(output_dir, 0o444) # Read-only
156+
157+
config = he_as.AssemblerRunConfig(
158+
input_file=str(input_file),
159+
input_mem_file=str(mem_file),
160+
output_dir=str(output_dir),
161+
)
162+
163+
with pytest.raises(Exception, match="Failed to write to output location"):
164+
he_as.main(config, verbose=False)
165+
166+
167+
def test_main_sets_global_config(tmp_path, monkeypatch):
168+
"""Test that main correctly sets GlobalConfig values."""
169+
input_file = tmp_path / "kernel.tw"
170+
input_file.write_text("dummy")
171+
mem_file = tmp_path / "kernel.mem"
172+
mem_file.write_text("mem")
173+
174+
config = he_as.AssemblerRunConfig(
175+
input_file=str(input_file),
176+
input_mem_file=str(mem_file),
177+
has_hbm=False,
178+
use_xinstfetch=True,
179+
suppress_comments=True,
180+
debug_verbose=2,
181+
)
182+
183+
asm_mock = mock.Mock(return_value=(1, 2, 3, 4, 5))
184+
monkeypatch.setattr(he_as, "asmisaAssemble", asm_mock)
185+
186+
he_as.main(config, verbose=False)
187+
188+
assert he_as.GlobalConfig.hasHBM is False
189+
assert he_as.GlobalConfig.useXInstFetch is True
190+
assert he_as.GlobalConfig.suppress_comments is True
191+
assert he_as.GlobalConfig.debugVerbose == 2
192+
193+
194+
def test_main_verbose_output(tmp_path, monkeypatch, capsys):
195+
"""Test that main prints verbose output when enabled."""
196+
input_file = tmp_path / "kernel.tw"
197+
input_file.write_text("dummy")
198+
mem_file = tmp_path / "kernel.mem"
199+
mem_file.write_text("mem")
200+
201+
config = he_as.AssemblerRunConfig(
202+
input_file=str(input_file),
203+
input_mem_file=str(mem_file),
204+
)
205+
206+
asm_mock = mock.Mock(return_value=(10, 2, 5, 1.5, 2.5))
207+
monkeypatch.setattr(he_as, "asmisaAssemble", asm_mock)
208+
209+
he_as.main(config, verbose=True)
210+
211+
captured = capsys.readouterr()
212+
assert "Output:" in captured.out
213+
assert "Total XInstructions: 10" in captured.out
214+
assert "Deps time: 1.5" in captured.out
215+
assert "Scheduling time: 2.5" in captured.out
216+
assert "Minimum idle cycles: 5" in captured.out
217+
assert "Minimum nops required: 2" in captured.out
218+
219+
220+
def test_config_as_dict(tmp_path):
221+
"""Test that AssemblerRunConfig.as_dict returns all config values."""
222+
input_file = tmp_path / "kernel.tw"
223+
input_file.write_text("")
224+
225+
config = he_as.AssemblerRunConfig(
226+
input_file=str(input_file),
227+
has_hbm=False,
228+
hbm_size=128,
229+
)
230+
231+
config_dict = config.as_dict()
232+
assert "input_file" in config_dict
233+
assert "has_hbm" in config_dict
234+
assert config_dict["has_hbm"] is False
235+
assert config_dict["hbm_size"] == 128
236+
237+
238+
def test_config_str_representation(tmp_path):
239+
"""Test that AssemblerRunConfig has a string representation."""
240+
input_file = tmp_path / "kernel.tw"
241+
input_file.write_text("")
242+
243+
config = he_as.AssemblerRunConfig(input_file=str(input_file))
244+
config_str = str(config)
245+
246+
assert "input_file" in config_str
247+
assert str(input_file) in config_str

0 commit comments

Comments
 (0)