mirror of
https://github.com/zama-ai/concrete.git
synced 2026-01-13 23:08:14 -05:00
103 lines
3.0 KiB
Python
103 lines
3.0 KiB
Python
"""
|
|
Tests of execution of the progress bar.
|
|
"""
|
|
|
|
from concrete import fhe
|
|
|
|
STEP_ESCAPED = "\\E2\\96\\88"
|
|
|
|
|
|
def test_progress_bar(helpers, monkeypatch):
|
|
"""
|
|
Test progress bar with interactive terminal.
|
|
"""
|
|
monkeypatch.setattr(
|
|
"concrete.fhe.mlir.converter.Converter.stdout_with_ansi_support", lambda: True
|
|
)
|
|
|
|
def function(x):
|
|
acc = x
|
|
with fhe.tag("init"):
|
|
acc += x
|
|
with fhe.tag("loop"):
|
|
for i in range(100):
|
|
with fhe.tag(f"round_{i}"):
|
|
acc += x
|
|
for _ in range(100):
|
|
acc += x
|
|
return acc
|
|
|
|
configuration = helpers.configuration().fork(
|
|
show_progress=True,
|
|
progress_title="TestTitle: ",
|
|
progress_tag=1,
|
|
)
|
|
compiler = fhe.Compiler(function, {"x": "encrypted"})
|
|
inputset = [0, 4]
|
|
circuit = compiler.compile(inputset, configuration)
|
|
next_is_tracing = True
|
|
percents = [f" {i:>3}% " for i in range(101)]
|
|
steps_done = [i // 2 for i in range(101)]
|
|
tags = ["init"] + ["(loop)"] * 49 + [""] * 51
|
|
for line in circuit.mlir.splitlines():
|
|
if "FHE.add_eint" in line:
|
|
next_is_tracing = True
|
|
elif "Tracing.trace_message" in line:
|
|
assert next_is_tracing
|
|
msg = line.split("msg = ")[1].split('"')[1]
|
|
assert configuration.progress_title in msg
|
|
assert percents.pop(0) in msg
|
|
assert msg.count(STEP_ESCAPED) == steps_done.pop(0)
|
|
assert tags.pop(0) in msg.rsplit("%", 1)[1]
|
|
next_is_tracing = False
|
|
|
|
|
|
def test_progress_bar_no_ansi(helpers, monkeypatch):
|
|
"""
|
|
Test progress bar with non-interactive terminal (pipe, notebook, etc).
|
|
"""
|
|
monkeypatch.setattr(
|
|
"concrete.fhe.mlir.converter.Converter.stdout_with_ansi_support", lambda: False
|
|
)
|
|
|
|
def function(x):
|
|
acc = x
|
|
for _ in range(100):
|
|
acc += x
|
|
return acc
|
|
|
|
configuration = helpers.configuration().fork(show_progress=True)
|
|
compiler = fhe.Compiler(function, {"x": "encrypted"})
|
|
|
|
inputset = [0, 4]
|
|
circuit = compiler.compile(inputset, configuration)
|
|
next_is_tracing = True
|
|
for line in circuit.mlir.splitlines():
|
|
print(line)
|
|
if "FHE.add_eint" in line:
|
|
next_is_tracing = True
|
|
elif "Tracing.trace_message" in line:
|
|
if "______" in line:
|
|
continue
|
|
assert next_is_tracing
|
|
msg = line.split("msg = ")[1].split('"')[1]
|
|
assert msg.count(STEP_ESCAPED) == 1 # add one stone
|
|
next_is_tracing = False
|
|
assert "100%" in msg
|
|
|
|
|
|
def test_progress_bar_empty_circuit(helpers):
|
|
"""
|
|
Test progress bar on empty circuit.
|
|
"""
|
|
|
|
def function(x):
|
|
return x
|
|
|
|
configuration = helpers.configuration().fork(show_progress=True)
|
|
compiler = fhe.Compiler(function, {"x": "encrypted"})
|
|
|
|
inputset = [0, 4]
|
|
circuit = compiler.compile(inputset, configuration)
|
|
assert "Tracing.trace_message" not in circuit.mlir
|