1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
import dataclasses
import logging
import platform
import shutil
from pathlib import Path

from testlib.environ import environ

import pytest

SLASHES_IN_STORE_PATH_UNTIL_PACKAGE = "/nix/store/hash-program_name/".count("/")


logger = logging.getLogger(__name__)


@dataclasses.dataclass
class _ManagedPath:
    """
    Wrapper class to handle building the `PATH` environment variable
    """

    build_shell: dataclasses.InitVar[str | None]
    build_env: dataclasses.InitVar[str | None] = None
    """statically linked shell to use within builds which provides coreutils functionality"""
    _path: list[str] = dataclasses.field(default_factory=list)

    def __post_init__(self, build_shell: str | None, build_env: str | None):
        if build_shell:
            self.prepend(build_shell)
        if build_env:
            for part in build_env.split(":"):
                self.append(part)

    def to_path(self) -> str:
        """
        :return: string to be put into the `PATH` environment variable containing all added paths/programs
        """
        return ":".join(self._path)

    def prepend(self, exec_path: str | Path) -> "_ManagedPath":
        """
        Adds the given file or folder at the FRONT of the path variable
        :param exec_path: executable or folder containing executables to be added
        :return: self, to allow for chaining
        """
        self._path.insert(0, str(exec_path))
        return self

    def append(self, exec_path: str | Path) -> "_ManagedPath":
        """
        Adds the given file or folder at the END of the path variable
        :param exec_path: executable or folder containing executables to be added
        :return: self, to allow for chaining
        """
        self._path.append(str(exec_path))
        return self

    def insert_at(self, exec_path: str | Path, index: int) -> "_ManagedPath":
        """
        Adds the given file or folder at the GIVEN INDEX of the path variable
        :param exec_path: executable or folder containing executable to be added
        :param index: where to insert the path
        :return: self, to allow for chaining
        """
        self._path.insert(index, str(exec_path))
        return self

    def add_program(self, program_name: str, all_associated: bool = True) -> "_ManagedPath":
        """
        Adds the given program to the path by name.
        :param program_name: executable/program to add
        :param all_associated: if True, the folder containing the executable will be added instead. Otherwise, only the provided executable will be added
        :raises ValueError: if the program could not be found
        :return: self, to allow for chaining
        """
        path = shutil.which(program_name)
        if path is None:
            msg = f"Couldn't find program {program_name!r}"
            raise ValueError(msg)
        # Convert to path object for better checking and operations
        path = Path(path)
        if all_associated and not path.is_dir():
            path = path.parent
        # handle as string within the data structure
        path = str(path)
        if path not in self._path:
            self._path.append(path)
        return self

    def remove_path(self, exec_path: str | Path) -> "_ManagedPath":
        """
        Removes the given file or folder from the path
        :param exec_path: file or folder to remove
        :raises ValueError: if the file or folder could not be found
        :return: self, to allow for chaining
        """
        self._path.remove(str(exec_path))
        return self

    def remove_program(self, program_name: str) -> "_ManagedPath":
        """
        Removes the given executable/program from the path by name
        :param program_name: executable/program to remove
        :raises ValueError: if the program was not found or isn't present in path
        :return: self, to allow for chaining
        """
        path = shutil.which(program_name)
        if path is None:
            msg = f"Couldn't find program {program_name!r}"
            raise ValueError(msg)
        if path in self._path or (path := str(Path(path).parent)) in self._path:
            self.remove_path(path)
        else:
            # Mirror behavior of `remove_path`
            msg = f"path.remove({program_name}): {program_name} not in path"
            raise ValueError(msg)
        return self

    def to_sandbox_paths(self) -> list[str]:
        """
        :return: list of strings to put into the `sandbox_paths` nix setting
        """
        ret = []
        for p in self._path:
            if p.startswith("/nix/store/"):
                # adds the entire package to the sandbox,
                # to ensure that dependencies and libraries from within the package are also present
                ret.append("/".join(p.split("/")[:SLASHES_IN_STORE_PATH_UNTIL_PACKAGE]))
            else:
                ret.append(p)
        return ret

    def which(self, program_name: str) -> Path:
        path = shutil.which(program_name, path=self.to_path())
        if not path:
            raise ValueError(f"{program_name} is not in configured path")
        return Path(path)


@dataclasses.dataclass
class _Dirs:
    test_root: Path
    home: Path
    nix_log_dir: Path | None
    nix_state_dir: Path | None
    nix_conf_dir: Path | None
    nix_bin_dir: Path | None
    # this one *must not* be NIX_STORE_DIR, otherwise lix will pick it up
    # and misconfigure itself. the config system is unbelievable bullshit
    real_store_dir: Path | None
    cache_dir: Path | None
    xdg_cache_home: Path | None
    tmpdir: Path | None
    """used for nar caching"""
    nix_store_dir: Path | None = None

    def get_env_keys(self) -> set[str]:
        return {f.name.upper() for f in dataclasses.fields(self)}

    def to_env_vars(self) -> dict[str, str]:
        return {k.upper(): v for k, v in dataclasses.asdict(self).items() if v is not None}


class ManagedEnv:
    def __init__(self, tmp_path: Path):
        # Things fetched from the global env
        build_shell = environ.get("BUILD_TEST_SHELL")
        global_path = environ.get("PATH")
        build_env = environ.get("BUILD_TEST_ENV")
        # `NIX_BIN_DIR` either propagated from us or set by meson
        # Set to the codebase internal output if started standalone
        # This is where the current lix binaries are located.
        # local import to avoid cyclic dependencies
        from testlib.utils import lix_base_folder  # noqa: PLC0415

        lix_bin = Path(environ.get("NIX_BIN_DIR", lix_base_folder / "outputs/out/bin"))

        self._env = {}
        self.path = _ManagedPath(build_shell, build_env)
        self._tmp_path = tmp_path
        self.shell_dir = build_shell or "/bin"
        self.build_env = build_env

        self.dirs = _Dirs(
            test_root=self._get_dir(""),
            home=self._get_dir("test-home"),
            nix_log_dir=self._get_dir("var/log/nix"),
            nix_state_dir=self._get_dir("nix/var/nix"),
            nix_conf_dir=self._get_dir("etc/nix"),
            nix_bin_dir=lix_bin,
            real_store_dir=self._get_dir("nix/store"),
            nix_store_dir=self._get_dir("nix/store"),
            cache_dir=self._get_dir("test-binary-cache"),
            xdg_cache_home=self._get_dir("test-home/.cache"),
            tmpdir=self._get_dir("tmp"),
        )
        self.path.prepend(self.dirs.nix_bin_dir)
        self.path.prepend(self.shell_dir)
        self.init_defaults(global_path)

        self._registered_commands = []

    def _get_dir(self, sub_path: str) -> Path:
        p = self._tmp_path / sub_path
        p.mkdir(parents=True, exist_ok=True)
        return p

    def init_defaults(self, global_path: str):
        self._env = {
            # Do not use the system-wide or local config for git, but *none* instead
            "GIT_CONFIG_SYSTEM": "/dev/null",
            # Shell to use, required by lix to run sub processes / commands
            "SHELL": f"{self.shell_dir}/sh",
            # when writing things to the terminal (esp with man pages) use cat, to print the full output to stdout
            "PAGER": "cat",
            "BUILD_TEST_SHELL": self.shell_dir,
        }
        if self.build_env:
            self._env["BUILD_TEST_ENV"] = self.build_env
        if platform.system() == "Darwin":
            # Darwin / Apple behaves differently and requires _NIX_TEST_NO_SANDBOX to be set for whatever reason
            self._env |= {"_NIX_TEST_NO_SANDBOX": "1"}
            # Copy coreutils from the global path to maintain availability of commands that are not part of
            # XCode Developer Tools and provided by busybox on Linux, which does not build on Darwin
            for p in global_path.split(":"):
                if Path(p).exists() and "coreutils" in p:
                    self.path.append(p)
                    break

    def set_env(self, name: str, value: str):
        if name in self.dirs.get_env_keys():
            msg = f"Overriding paths should be done using the `env.dirs` attribute, use `env.dirs.{name.lower()}` instead."
            raise ValueError(msg)
        if name == "PATH":
            msg = "Setting of path not supported. use `env.path` instead"
            raise ValueError(msg)
        if value is None:
            msg = "setting to `None` is not allowed. did you mean to use `env.unset_env`?"
            raise ValueError(msg)
        self._env[name] = value

    def __setitem__(self, key: str, value: str) -> None:
        return self.set_env(key, value)

    def get_env(self, name: str, default: str | Path | None = None) -> str | Path | None:
        if name in self.dirs.get_env_keys():
            return getattr(self.dirs, name.lower())
        if name == "PATH":
            msg = "getting of path not supported, use `env.path` instead"
            raise ValueError(msg)
        return self._env.get(name, default)

    def __getitem__(self, item: str) -> str | Path | None:
        itm = self.get_env(item)
        if itm is None:
            msg = f"{itm} is not set"
            raise KeyError(msg)
        return itm

    def unset_env(self, name: str) -> str | None:
        if name in self.dirs.get_env_keys():
            msg = f"Overriding paths should be done using the `env.dirs` attribute, use `env.dirs.{name.lower()}` instead."
            raise ValueError(msg)
        return self._env.pop(name, None)

    def __delitem__(self, key: str) -> str | None:
        itm = self.unset_env(key)
        if itm is None:
            msg = f"{itm} is not set"
            raise KeyError(msg)
        return itm

    def to_env(self) -> dict[str, str]:
        ret = self.dirs.to_env_vars()
        ret["PATH"] = self.path.to_path()
        for k, v in self._env.copy().items():
            if v is not None:
                ret[k] = v
            else:
                logger.warning("environment variable '%s' is none", k)
        return ret


@pytest.fixture
def env(tmp_path: Path) -> ManagedEnv:
    env_ = ManagedEnv(tmp_path)
    yield env_
    for cmd in env_._registered_commands:
        if not cmd._run:
            e = ValueError("Command has not been run")
            e.add_note(str(cmd))
            raise e