Added test for onetrace_enabled() context manager. Fixed docstring example

Author: oleksandr-pavlyk

dpctl/tests/test_utils.py

@@ -111,3 +111,13 @@ def validate_usm_type_arg():
         dpctl.utils.validate_usm_type("inv", allow_none=True)
     with pytest.raises(ValueError):
         dpctl.utils.validate_usm_type("inv", allow_none=False)
+
+
+def test_onetrace_enabled():
+    import os
+
+    v_name = "PTI_ENABLE_COLLECTION"
+    v_v = os.getenv(v_name, None)
+    with dpctl.utils.onetrace_enabled():
+        assert os.getenv(v_name, None) == "1"
+    assert os.getenv(v_name, None) == v_v

dpctl/utils/_onetrace_context.py

@@ -30,7 +30,7 @@ def onetrace_enabled():
             $ onetrace --conditional-collection -v -t --demangle \
                 python script.py
 
-        .. Example:: bash
+        .. Example:: python
 
             import dpctl.tensor as dpt
             from dpctl.utils import onetrace_enabled