Repository: stan-dev/cmdstanpy

Issue:
Setting CMDSTAN environment variable causes test to fail
#### Summary:
cmdstanpy properly finds cmdstan when `CMDSTAN` environment is set, but one test does fail. 

#### Description:
There is a test to see that cmdstan is located at `~/.cmdstanpy/cmdstan-<version>`, but a user can change the location of cmdstan by setting `CMDSTAN` environment variable.

https://github.com/stan-dev/cmdstanpy/blob/f20c9bb9002387dfc662fe84d98d24f9df1850e8/test/test_utils.py#L35-L40

#### Current Version:
master


Comments:

Patch:
diff --git a/cmdstanpy/__init__.py b/cmdstanpy/__init__.py
index 9f1dbbb..45dc8e6 100644
--- a/cmdstanpy/__init__.py
+++ b/cmdstanpy/__init__.py
@@ -5,7 +5,7 @@ import atexit
 import shutil
 import tempfile
 
-STANSUMMARY_STATS = [
+_STANSUMMARY_STATS = [
     'Mean',
     'MCSE',
     'StdDev',
@@ -17,20 +17,24 @@ STANSUMMARY_STATS = [
     'R_hat',
 ]
 
-TMPDIR = tempfile.mkdtemp()
+_TMPDIR = tempfile.mkdtemp()
 
 
-def cleanup_tmpdir():
-    """Force deletion of TMPDIR."""
-    print('deleting tmpfiles dir: {}'.format(TMPDIR))
-    shutil.rmtree(TMPDIR, ignore_errors=True)
+def _cleanup_tmpdir():
+    """Force deletion of _TMPDIR."""
+    print('deleting tmpfiles dir: {}'.format(_TMPDIR))
+    shutil.rmtree(_TMPDIR, ignore_errors=True)
     print('done')
 
 
-atexit.register(cleanup_tmpdir)
+atexit.register(_cleanup_tmpdir)
 
 
 from .utils import set_cmdstan_path, cmdstan_path, set_make_env, install_cmdstan # noqa
 from .stanfit import CmdStanMCMC, CmdStanMLE, CmdStanGQ, CmdStanVB # noqa
 from .model import CmdStanModel # noqa
 from ._version import __version__ # noqa
+
+__all__ = ['set_cmdstan_path', 'cmdstan_path', 'set_make_env',
+            'install_cmdstan', 'CmdStanMCMC', 'CmdStanMLE',
+            'CmdStanGQ', 'CmdStanVB', 'CmdStanModel']
diff --git a/cmdstanpy/stanfit.py b/cmdstanpy/stanfit.py
index 9ca97d5..33cb70e 100644
--- a/cmdstanpy/stanfit.py
+++ b/cmdstanpy/stanfit.py
@@ -11,7 +11,7 @@ from time import time
 import numpy as np
 import pandas as pd
 
-from cmdstanpy import TMPDIR
+from cmdstanpy import _TMPDIR
 from cmdstanpy.utils import (
     check_sampler_csv,
     scan_optimize_csv,
@@ -54,7 +54,7 @@ class RunSet:
         if args.output_dir is not None:
             output_dir = args.output_dir
         else:
-            output_dir = TMPDIR
+            output_dir = _TMPDIR
 
         self._csv_files = []
         self._diagnostic_files = [None for _ in range(chains)]
@@ -77,7 +77,7 @@ class RunSet:
             if args.save_diagnostics:
                 if args.output_dir is None:
                     diag_file = create_named_text_file(
-                        dir=TMPDIR,
+                        dir=_TMPDIR,
                         prefix='{}-diagnostic-{}-'.format(file_basename, i + 1),
                         suffix='.csv',
                     )
@@ -202,7 +202,7 @@ class RunSet:
                 )
 
             path, filename = os.path.split(self._csv_files[i])
-            if path == TMPDIR:  # cleanup tmpstr in filename
+            if path == _TMPDIR:  # cleanup tmpstr in filename
                 root, ext = os.path.splitext(filename)
                 rlist = root.split('-')
                 root = '-'.join(rlist[:-1])
@@ -431,7 +431,7 @@ class CmdStanMCMC:
             self.runset._args.model_name, self.runset.chains
         )
         tmp_csv_path = create_named_text_file(
-            dir=TMPDIR, prefix=tmp_csv_file, suffix='.csv'
+            dir=_TMPDIR, prefix=tmp_csv_file, suffix='.csv'
         )
         cmd = [
             cmd_path,
diff --git a/cmdstanpy/utils.py b/cmdstanpy/utils.py
index 5f80dac..e7dd555 100644
--- a/cmdstanpy/utils.py
+++ b/cmdstanpy/utils.py
@@ -24,7 +24,7 @@ import numpy as np
 import pandas as pd
 
 
-from cmdstanpy import TMPDIR
+from cmdstanpy import _TMPDIR
 
 EXTENSION = '.exe' if platform.system() == 'Windows' else ''
 
@@ -66,7 +66,7 @@ class MaybeDictToFilePath():
         for obj in objs:
             if isinstance(obj, dict):
                 data_file = create_named_text_file(
-                    dir=TMPDIR, prefix='', suffix='.json'
+                    dir=_TMPDIR, prefix='', suffix='.json'
                 )
                 self._logger.debug('input tempfile: %s', data_file)
                 if any(
@@ -192,6 +192,7 @@ def cmdstan_path() -> str:
                 'run command line script "install_cmdstan"'
             )
         cmdstan = os.path.join(cmdstan_dir, latest_cmdstan)
+        os.environ['CMDSTAN'] = cmdstan
     validate_cmdstan_path(cmdstan)
     return cmdstan
 


Test patch:
diff --git a/test/test_cmdstan_args.py b/test/test_cmdstan_args.py
index c709378..cf6f111 100644
--- a/test/test_cmdstan_args.py
+++ b/test/test_cmdstan_args.py
@@ -4,7 +4,7 @@ import os
 import platform
 import unittest
 
-from cmdstanpy import TMPDIR
+from cmdstanpy import _TMPDIR
 from cmdstanpy.cmdstan_args import (
     Method,
     SamplerArgs,
@@ -496,7 +496,7 @@ class CmdStanArgsTest(unittest.TestCase):
         # TODO: read-only dir test for Windows - set ACLs, not mode
         if platform.system() == 'Darwin' or platform.system() == 'Linux':
             with self.assertRaises(ValueError):
-                read_only = os.path.join(TMPDIR, 'read_only')
+                read_only = os.path.join(_TMPDIR, 'read_only')
                 os.mkdir(read_only, mode=0o444)
                 CmdStanArgs(
                     model_name='bernoulli',
diff --git a/test/test_utils.py b/test/test_utils.py
index d319577..0b72a0f 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -10,7 +10,7 @@ import random
 
 import numpy as np
 
-from cmdstanpy import TMPDIR
+from cmdstanpy import _TMPDIR
 from cmdstanpy.utils import (
     cmdstan_path,
     set_cmdstan_path,
@@ -34,19 +34,44 @@ DATAFILES_PATH = os.path.join(HERE, 'data')
 
 class CmdStanPathTest(unittest.TestCase):
     def test_default_path(self):
-        abs_rel_path = os.path.expanduser(
-            os.path.join('~', '.cmdstanpy', 'cmdstan')
-        )
-        self.assertTrue(cmdstan_path().startswith(abs_rel_path))
+        cur_value = None
+        if 'CMDSTAN' in os.environ:
+            cur_value = os.environ['CMDSTAN']
+        try:
+            if 'CMDSTAN' in os.environ:
+                self.assertEqual(cmdstan_path(), os.environ['CMDSTAN'])
+                path = os.environ['CMDSTAN']
+                del os.environ['CMDSTAN']
+                self.assertFalse('CMDSTAN' in os.environ)
+                set_cmdstan_path(path)
+                self.assertEqual(cmdstan_path(), path)
+                self.assertTrue('CMDSTAN' in os.environ)
+            else:
+                install_dir = os.path.expanduser(
+                    os.path.join('~', '.cmdstanpy')
+                )
+                install_version = os.path.expanduser(
+                    os.path.join(install_dir, get_latest_cmdstan(install_dir))
+                )
+                self.assertTrue(
+                    os.path.samefile(cmdstan_path(), install_version)
+                )
+                self.assertTrue('CMDSTAN' in os.environ)
+        finally:
+            if cur_value is not None:
+                os.environ['CMDSTAN'] = cur_value
+            else:
+                if 'CMDSTAN' in os.environ:
+                    del os.environ['CMDSTAN']
 
     def test_non_spaces_location(self):
-        good_path = os.path.join(TMPDIR, 'good_dir')
+        good_path = os.path.join(_TMPDIR, 'good_dir')
         with TemporaryCopiedFile(good_path) as (pth, is_changed):
             self.assertEqual(pth, good_path)
             self.assertFalse(is_changed)
 
         # prepare files for test
-        bad_path = os.path.join(TMPDIR, 'bad dir')
+        bad_path = os.path.join(_TMPDIR, 'bad dir')
         os.makedirs(bad_path, exist_ok=True)
         stan = os.path.join(DATAFILES_PATH, 'bernoulli.stan')
         stan_bad = os.path.join(bad_path, 'bad name.stan')
@@ -70,12 +95,16 @@ class CmdStanPathTest(unittest.TestCase):
         shutil.rmtree(bad_path, ignore_errors=True)
 
     def test_set_path(self):
-        install_dir = os.path.expanduser(os.path.join('~', '.cmdstanpy'))
-        install_version = os.path.expanduser(
-            os.path.join(install_dir, get_latest_cmdstan(install_dir))
-        )
-        set_cmdstan_path(install_version)
-        self.assertEqual(install_version, cmdstan_path())
+        if 'CMDSTAN' in os.environ:
+            self.assertEqual(cmdstan_path(), os.environ['CMDSTAN'])
+        else:
+            install_dir = os.path.expanduser(os.path.join('~', '.cmdstanpy'))
+            install_version = os.path.expanduser(
+                os.path.join(install_dir, get_latest_cmdstan(install_dir))
+            )
+            set_cmdstan_path(install_version)
+            self.assertEqual(install_version, cmdstan_path())
+            self.assertEqual(install_version, os.environ['CMDSTAN'])
 
     def test_validate_path(self):
         install_dir = os.path.expanduser(os.path.join('~', '.cmdstanpy'))
@@ -119,32 +148,32 @@ class CmdStanPathTest(unittest.TestCase):
 
     def test_jsondump(self):
         dict_list = {'a': [1.0, 2.0, 3.0]}
-        file_list = os.path.join(TMPDIR, 'list.json')
+        file_list = os.path.join(_TMPDIR, 'list.json')
         jsondump(file_list, dict_list)
         with open(file_list) as fd:
             self.assertEqual(json.load(fd), dict_list)
 
         dict_vec = {'a': np.repeat(3, 4)}
-        file_vec = os.path.join(TMPDIR, 'vec.json')
+        file_vec = os.path.join(_TMPDIR, 'vec.json')
         jsondump(file_vec, dict_vec)
         with open(file_vec) as fd:
             self.assertEqual(json.load(fd), dict_vec)
 
         dict_zero_vec = {'a': []}
-        file_zero_vec = os.path.join(TMPDIR, 'empty_vec.json')
+        file_zero_vec = os.path.join(_TMPDIR, 'empty_vec.json')
         jsondump(file_zero_vec, dict_zero_vec)
         with open(file_zero_vec) as fd:
             self.assertEqual(json.load(fd), dict_zero_vec)
 
         dict_zero_matrix = {'a': [[], [], []]}
-        file_zero_matrix = os.path.join(TMPDIR, 'empty_matrix.json')
+        file_zero_matrix = os.path.join(_TMPDIR, 'empty_matrix.json')
         jsondump(file_zero_matrix, dict_zero_matrix)
         with open(file_zero_matrix) as fd:
             self.assertEqual(json.load(fd), dict_zero_matrix)
 
         arr = np.zeros(shape=(5, 0))
         dict_zero_matrix = {'a': arr}
-        file_zero_matrix = os.path.join(TMPDIR, 'empty_matrix.json')
+        file_zero_matrix = os.path.join(_TMPDIR, 'empty_matrix.json')
         jsondump(file_zero_matrix, dict_zero_matrix)
         with open(file_zero_matrix) as fd:
             self.assertEqual(json.load(fd), dict_zero_matrix)
@@ -282,7 +311,7 @@ class WindowsShortPath(unittest.TestCase):
     def test_windows_short_path_directory(self):
         if platform.system() != 'Windows':
             return
-        original_path = os.path.join(TMPDIR, 'new path')
+        original_path = os.path.join(_TMPDIR, 'new path')
         os.makedirs(original_path, exist_ok=True)
         assert os.path.exists(original_path)
         assert ' ' in original_path
@@ -294,7 +323,7 @@ class WindowsShortPath(unittest.TestCase):
     def test_windows_short_path_file(self):
         if platform.system() != 'Windows':
             return
-        original_path = os.path.join(TMPDIR, 'new path', 'my_file.csv')
+        original_path = os.path.join(_TMPDIR, 'new path', 'my_file.csv')
         os.makedirs(os.path.split(original_path)[0], exist_ok=True)
         assert os.path.exists(os.path.split(original_path)[0])
         assert ' ' in original_path
@@ -309,7 +338,7 @@ class WindowsShortPath(unittest.TestCase):
         """Test that the function doesn't touch filename."""
         if platform.system() != 'Windows':
             return
-        original_path = os.path.join(TMPDIR, 'new path', 'my file.csv')
+        original_path = os.path.join(_TMPDIR, 'new path', 'my file.csv')
         os.makedirs(os.path.split(original_path)[0], exist_ok=True)
         assert os.path.exists(os.path.split(original_path)[0])
         assert ' ' in original_path