code: def test_set_omp_threads():
      if have_openmp:
          # set threads to default
          _restore_omp_threads()
          assert_equal(thread_count(), default_threads)
  
          # change number of threads
          nthreads_new = default_threads + 1
          _set_omp_threads(nthreads_new)
          assert_equal(thread_count(), nthreads_new)
  
          # restore back to default
          _restore_omp_threads()
          assert_equal(thread_count(), default_threads)
      else:
          assert_equal(thread_count(), 1)
          assert_equal(cpu_count(), 1)
  
file_path: dipy/dipy/utils/tests/test_omp.py
imports: import os
  from dipy.utils.omp import cpu_count, thread_count, default_threads, _set_omp_threads, _restore_omp_threads, have_openmp, determine_num_threads
  import pytest
  from dipy.utils.parallel import has_ray, has_joblib
  from numpy.testing import assert_equal, assert_raises
label_desc: Use image node_id 1 for calling a global function with example usage: test_set_omp_threads() without return types