1#!/usr/bin/env python3
2# Copyright 2023 The Bazel Authors. All rights reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16
17import os
18import platform
19import subprocess
20import sys
21import unittest
22from pathlib import Path
23
24from rules_python.python.runfiles import runfiles
25
26
27class PipRepositoryAnnotationsTest(unittest.TestCase):
28    maxDiff = None
29
30    def wheel_pkg_dir(self) -> str:
31        env = os.environ.get("WHEEL_PKG_DIR")
32        self.assertIsNotNone(env)
33        return env
34
35    def test_build_content_and_data(self):
36        r = runfiles.Create()
37        rpath = r.Rlocation(
38            "pip_repository_annotations_example/external/{}/generated_file.txt".format(
39                self.wheel_pkg_dir()
40            )
41        )
42        generated_file = Path(rpath)
43        self.assertTrue(generated_file.exists())
44
45        content = generated_file.read_text().rstrip()
46        self.assertEqual(content, "Hello world from build content file")
47
48    def test_copy_files(self):
49        r = runfiles.Create()
50        rpath = r.Rlocation(
51            "pip_repository_annotations_example/external/{}/copied_content/file.txt".format(
52                self.wheel_pkg_dir()
53            )
54        )
55        copied_file = Path(rpath)
56        self.assertTrue(copied_file.exists())
57
58        content = copied_file.read_text().rstrip()
59        self.assertEqual(content, "Hello world from copied file")
60
61    def test_copy_executables(self):
62        r = runfiles.Create()
63        rpath = r.Rlocation(
64            "pip_repository_annotations_example/external/{}/copied_content/executable{}".format(
65                self.wheel_pkg_dir(),
66                ".exe" if platform.system() == "windows" else ".py",
67            )
68        )
69        executable = Path(rpath)
70        self.assertTrue(executable.exists())
71
72        proc = subprocess.run(
73            [sys.executable, str(executable)],
74            check=True,
75            stdout=subprocess.PIPE,
76            stderr=subprocess.PIPE,
77        )
78        stdout = proc.stdout.decode("utf-8").strip()
79        self.assertEqual(stdout, "Hello world from copied executable")
80
81    def test_data_exclude_glob(self):
82        current_wheel_version = "0.38.4"
83
84        r = runfiles.Create()
85        dist_info_dir = "pip_repository_annotations_example/external/{}/site-packages/wheel-{}.dist-info".format(
86            self.wheel_pkg_dir(),
87            current_wheel_version,
88        )
89
90        # Note: `METADATA` is important as it's consumed by https://docs.python.org/3/library/importlib.metadata.html
91        # `METADATA` is expected to be there to show dist-info files are included in the runfiles.
92        metadata_path = r.Rlocation("{}/METADATA".format(dist_info_dir))
93
94        # However, `WHEEL` was explicitly excluded, so it should be missing
95        wheel_path = r.Rlocation("{}/WHEEL".format(dist_info_dir))
96
97        # Because windows does not have `--enable_runfiles` on by default, the
98        # `runfiles.Rlocation` results will be different on this platform vs
99        # unix platforms. See `@rules_python//python/runfiles` for more details.
100        if platform.system() == "Windows":
101            self.assertIsNotNone(metadata_path)
102            self.assertIsNone(wheel_path)
103        else:
104            self.assertTrue(Path(metadata_path).exists())
105            self.assertFalse(Path(wheel_path).exists())
106
107    def requests_pkg_dir(self) -> str:
108        env = os.environ.get("REQUESTS_PKG_DIR")
109        self.assertIsNotNone(env)
110        return env
111
112    def test_extra(self):
113        # This test verifies that annotations work correctly for pip packages with extras
114        # specified, in this case requests[security].
115        r = runfiles.Create()
116        rpath = r.Rlocation(
117            "pip_repository_annotations_example/external/{}/generated_file.txt".format(
118                self.requests_pkg_dir()
119            )
120        )
121        generated_file = Path(rpath)
122        self.assertTrue(generated_file.exists())
123
124        content = generated_file.read_text().rstrip()
125        self.assertEqual(content, "Hello world from requests")
126
127
128if __name__ == "__main__":
129    unittest.main()
130