Skip to content

Commit b5fc478

Browse files
Added test for onetrace_enabled() context manager. Fixed docstring example
1 parent f216df6 commit b5fc478

File tree

2 files changed

+11
-1
lines changed

2 files changed

+11
-1
lines changed

dpctl/tests/test_utils.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,3 +111,13 @@ def validate_usm_type_arg():
111111
dpctl.utils.validate_usm_type("inv", allow_none=True)
112112
with pytest.raises(ValueError):
113113
dpctl.utils.validate_usm_type("inv", allow_none=False)
114+
115+
116+
def test_onetrace_enabled():
117+
import os
118+
119+
v_name = "PTI_ENABLE_COLLECTION"
120+
v_v = os.getenv(v_name, None)
121+
with dpctl.utils.onetrace_enabled():
122+
assert os.getenv(v_name, None) == "1"
123+
assert os.getenv(v_name, None) == v_v

dpctl/utils/_onetrace_context.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ def onetrace_enabled():
3030
$ onetrace --conditional-collection -v -t --demangle \
3131
python script.py
3232
33-
.. Example:: bash
33+
.. Example:: python
3434
3535
import dpctl.tensor as dpt
3636
from dpctl.utils import onetrace_enabled

0 commit comments

Comments
 (0)