python-dotenv/tests/test_zip_imports.py at 9d6fd0704e3c5914c4c6c72faefcfc86e864adb1 · bbc2/python-dotenv · GitHub

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
import os
import posixpath
import subprocess
import sys
import textwrap
from typing import List
from unittest import mock
from zipfile import ZipFile


def walk_to_root(path: str):
    last_dir = None
    current_dir = path
    while last_dir != current_dir:
        yield current_dir
        parent_dir = posixpath.dirname(current_dir)
        last_dir, current_dir = current_dir, parent_dir


class FileToAdd:
    def __init__(self, content: str, path: str):
        self.content = content
        self.path = path


def setup_zipfile(path, files: List[FileToAdd]):
    zip_file_path = path / "test.zip"
    dirs_init_py_added_to = set()
    with ZipFile(zip_file_path, "w") as zipfile:
        for f in files:
            zipfile.writestr(data=f.content, zinfo_or_arcname=f.path)
            for dirname in walk_to_root(posixpath.dirname(f.path)):
                if dirname not in dirs_init_py_added_to:
                    init_path = posixpath.join(dirname, "__init__.py")
                    print(f"setup_zipfile: {init_path}")
                    zipfile.writestr(data="", zinfo_or_arcname=init_path)
                    dirs_init_py_added_to.add(dirname)
    return zip_file_path


@mock.patch.object(sys, "path", list(sys.path))
def test_load_dotenv_gracefully_handles_zip_imports_when_no_env_file(tmp_path):
    zip_file_path = setup_zipfile(
        tmp_path,
        [
            FileToAdd(
                content=textwrap.dedent(
                    """
            from dotenv import load_dotenv

            load_dotenv()
        """
                ),
                path="child1/child2/test.py",
            ),
        ],
    )

    # Should run without an error
    sys.path.append(str(zip_file_path))
    import child1.child2.test  # noqa


def test_load_dotenv_outside_zip_file_when_called_in_zipfile(tmp_path):
    zip_file_path = setup_zipfile(
        tmp_path,
        [
            FileToAdd(
                content=textwrap.dedent(
                    """
            from dotenv import load_dotenv

            load_dotenv()
        """
                ),
                path="child1/child2/test.py",
            ),
        ],
    )
    dotenv_path = tmp_path / ".env"
    dotenv_path.write_bytes(b"A=x")
    code_path = tmp_path / "code.py"
    code_path.write_text(
        textwrap.dedent(
            f"""
            import os
            import sys

            sys.path.append({str(zip_file_path)!r})

            import child1.child2.test

            print(os.environ['A'])
            """
        )
    )

    result = subprocess.run(
        [sys.executable, str(code_path)],
        capture_output=True,
        check=True,
        cwd=tmp_path,
        text=True,
        env={
            k: v for k, v in os.environ.items() if k.upper() != "A"
        },  # env without 'A'
    )

    assert result.stdout == "x\n"