src/test/py/bazel/bzlmod/test_utils.py - bazel - Git at Google

 # pylint: disable=invalid-name
 # pylint: disable=g-long-ternary
 # Copyright 2021 The Bazel Authors. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #    http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Test utils for Bzlmod."""

 import base64
 import functools
 import hashlib
 import http.server
 import json
 import os
 import pathlib
 import shutil
 import threading
 import urllib.request
 import zipfile


 def download(url):
   """Download a file and return its content in bytes."""
   response = urllib.request.urlopen(url)
   return response.read()


 def read(path):
   """Read a file and return its content in bytes."""
   with open(str(path), 'rb') as f:
     return f.read()


 def integrity(data):
   """Calculate the integration value of the data with sha256."""
   hash_value = hashlib.sha256(data)
   return 'sha256-' + base64.b64encode(hash_value.digest()).decode()


 def scratchFile(path, lines=None):
   """Creates a file at the given path with the given content."""
   with open(str(path), 'w') as f:
     if lines:
       for l in lines:
         f.write(l)
         f.write('\n')


 class Module:
   """A class to represent information of a Bazel module."""

   def __init__(self, name, version):
     self.name = name
     self.version = version
     self.archive_url = None
     self.strip_prefix = ''
     self.module_dot_bazel = None
     self.patches = []
     self.patch_strip = 0
     self.archive_type = None

   def set_source(self, archive_url, strip_prefix=None):
     self.archive_url = archive_url
     self.strip_prefix = strip_prefix
     return self

   def set_module_dot_bazel(self, module_dot_bazel):
     self.module_dot_bazel = module_dot_bazel
     return self

   def set_patches(self, patches, patch_strip):
     self.patches = patches
     self.patch_strip = patch_strip
     return self

   def set_archive_type(self, archive_type):
     self.archive_type = archive_type
     return self


 class BazelRegistry:
   """A class to help create a Bazel module project from scatch and add it into the registry."""

   def __init__(self, root, registry_suffix=''):
     self.root = pathlib.Path(root)
     self.projects = self.root.joinpath('projects')
     self.projects.mkdir(parents=True, exist_ok=True)
     self.archives = self.root.joinpath('archives')
     self.archives.mkdir(parents=True, exist_ok=True)
     self.registry_suffix = registry_suffix

   def setModuleBasePath(self, module_base_path):
     bazel_registry = {
         'module_base_path': module_base_path,
     }
     with self.root.joinpath('bazel_registry.json').open('w') as f:
       json.dump(bazel_registry, f, indent=4, sort_keys=True)

   def getURL(self):
     """Return the URL of this registry."""
     return self.root.resolve().as_uri()

   def generateCcSource(
       self,
       name,
       version,
       deps=None,
       repo_names=None,
       extra_module_file_contents=None,
   ):
     """Generate a cc project with given dependency information.

     1. The cc projects implements a hello_<lib_name> function.
     2. The hello_<lib_name> function calls the same function of its
     dependencies.
     3. The hello_<lib_name> function prints "<caller name> =>
     <lib_name@version>".
     4. The BUILD file references the dependencies as their desired repo names.

     Args:
       name:  The module name.
       version: The module version.
       deps: The dependencies of this module.
       repo_names: The desired repository name for some dependencies.
       extra_module_file_contents: Extra lines to append to the MODULE.bazel
         file.

     Returns:
       The generated source directory.
     """

     src_dir = self.projects.joinpath(name, version)
     src_dir.mkdir(parents=True, exist_ok=True)
     if not deps:
       deps = {}
     if not repo_names:
       repo_names = {}
     for dep in deps:
       if dep not in repo_names:
         repo_names[dep] = dep
     if not extra_module_file_contents:
       extra_module_file_contents = []

     def calc_repo_name_str(dep):
       if dep == repo_names[dep]:
         return ''
       return ', repo_name = "%s"' % repo_names[dep]

     scratchFile(src_dir.joinpath('WORKSPACE'))
     scratchFile(
         src_dir.joinpath('MODULE.bazel'),
         [
             'module(',
             '  name = "%s",' % name,
             '  version = "%s",' % version,
             '  compatibility_level = 1,',
             ')',
         ]
         + [
             'bazel_dep(name = "%s", version = "%s"%s)'
             % (dep, version, calc_repo_name_str(dep))
             for dep, version in deps.items()
         ]
         + extra_module_file_contents,
     )

     scratchFile(
         src_dir.joinpath(name.lower() + '.h'), [
             '#ifndef %s_H' % name.upper(),
             '#define %s_H' % name.upper(),
             '#include <string>',
             'void hello_%s(const std::string& caller);' % name.lower(),
             '#endif',
         ])
     scratchFile(
         src_dir.joinpath(name.lower() + '.cc'), [
             '#include <stdio.h>',
             '#include "%s.h"' % name.lower(),
         ] + ['#include "%s.h"' % dep.lower() for dep in deps] + [
             'void hello_%s(const std::string& caller) {' % name.lower(),
             '    std::string lib_name = "%s@%s%s";' %
             (name, version, self.registry_suffix),
             '    printf("%s => %s\\n", caller.c_str(), lib_name.c_str());',
         ] + ['    hello_%s(lib_name);' % dep.lower() for dep in deps] + [
             '}',
         ])
     scratchFile(
         src_dir.joinpath('BUILD'), [
             'package(default_visibility = ["//visibility:public"])',
             'cc_library(',
             '  name = "lib_%s",' % name.lower(),
             '  srcs = ["%s.cc"],' % name.lower(),
             '  hdrs = ["%s.h"],' % name.lower(),
         ] + ([
             '  deps = ["%s"],' % ('", "'.join([
                 '@%s//:lib_%s' % (repo_names[dep], dep.lower()) for dep in deps
             ])),
         ] if deps else []) + [
             ')',
         ])
     return src_dir

   def createArchive(self, name, version, src_dir, filename_pattern='%s.%s.zip'):
     """Create an archive with a given source directory."""
     zip_path = self.archives.joinpath(filename_pattern % (name, version))
     zip_obj = zipfile.ZipFile(str(zip_path), 'w')
     for foldername, _, filenames in os.walk(str(src_dir)):
       for filename in filenames:
         filepath = os.path.join(foldername, filename)
         zip_obj.write(filepath,
                       str(pathlib.Path(filepath).relative_to(src_dir)))
     zip_obj.close()
     return zip_path

   def addModule(self, module):
     """Add a module into the registry."""
     module_dir = self.root.joinpath('modules', module.name, module.version)
     module_dir.mkdir(parents=True, exist_ok=True)

     # Copy MODULE.bazel to the registry
     module_dot_bazel = module_dir.joinpath('MODULE.bazel')
     shutil.copy(str(module.module_dot_bazel), str(module_dot_bazel))

     # Create source.json & copy patch files to the registry
     source = {
         'url': module.archive_url,
         'integrity': integrity(download(module.archive_url)),
     }
     if module.strip_prefix:
       source['strip_prefix'] = module.strip_prefix

     if module.patches:
       patch_dir = module_dir.joinpath('patches')
       patch_dir.mkdir()
       source['patches'] = {}
       source['patch_strip'] = module.patch_strip
       for patch_path in module.patches:
         patch = pathlib.Path(patch_path)
         source['patches'][patch.name] = integrity(read(patch))
         shutil.copy(str(patch), str(patch_dir))

     if module.archive_type:
       source['archive_type'] = module.archive_type

     with module_dir.joinpath('source.json').open('w') as f:
       json.dump(source, f, indent=4, sort_keys=True)

   def createCcModule(
       self,
       name,
       version,
       deps=None,
       repo_names=None,
       patches=None,
       patch_strip=0,
       archive_pattern=None,
       archive_type=None,
       extra_module_file_contents=None,
   ):
     """Generate a cc project and add it as a module into the registry."""
     src_dir = self.generateCcSource(
         name, version, deps, repo_names, extra_module_file_contents
     )
     if archive_pattern:
       archive = self.createArchive(
           name, version, src_dir, filename_pattern=archive_pattern
       )
     else:
       archive = self.createArchive(name, version, src_dir)
     module = Module(name, version)
     module.set_source(archive.resolve().as_uri())
     module.set_module_dot_bazel(src_dir.joinpath('MODULE.bazel'))
     if patches:
       module.set_patches(patches, patch_strip)
     if archive_type:
       module.set_archive_type(archive_type)

     self.addModule(module)
     return self

   def addMetadata(self,
                   name,
                   homepage=None,
                   maintainers=None,
                   versions=None,
                   yanked_versions=None):
     """Generate a module metadata file and add it to the registry."""
     if maintainers is None:
       maintainers = []
     if versions is None:
       versions = []
     if yanked_versions is None:
       yanked_versions = {}

     module_dir = self.root.joinpath('modules', name)
     module_dir.mkdir(parents=True, exist_ok=True)

     metadata = {
         'homepage': homepage,
         'maintainers': maintainers,
         'versions': versions,
         'yanked_versions': yanked_versions
     }

     with module_dir.joinpath('metadata.json').open('w') as f:
       json.dump(metadata, f, indent=4, sort_keys=True)

     return self

   def createLocalPathModule(self, name, version, path, deps=None):
     """Add a local module into the registry."""
     module_dir = self.root.joinpath('modules', name, version)
     module_dir.mkdir(parents=True, exist_ok=True)

     # Create source.json & copy patch files to the registry
     source = {
         'type': 'local_path',
         'path': path,
     }

     if deps is None:
       deps = {}

     module_file_lines = [
         'module(',
         '  name = "%s",' % name,
         '  version = "%s",' % version,
         ')',
     ] + ['bazel_dep(name="%s",version="%s")' % p for p in deps.items()]
     scratchFile(module_dir.joinpath('MODULE.bazel'), module_file_lines)
     self.projects.joinpath(path).mkdir(exist_ok=True)
     scratchFile(self.projects.joinpath(path, 'MODULE.bazel'), module_file_lines)

     with module_dir.joinpath('source.json').open('w') as f:
       json.dump(source, f, indent=4, sort_keys=True)


 class StaticHTTPServer:
   """An HTTP server serving static files, optionally with authentication."""

   def __init__(self, root_directory, expected_auth=None):
     self.root_directory = root_directory
     self.expected_auth = expected_auth

   def __enter__(self):
     address = ('localhost', 0)  # assign random port
     handler = functools.partial(
         _Handler, self.root_directory, self.expected_auth
     )
     self.httpd = http.server.HTTPServer(address, handler)
     self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True)
     self.thread.start()
     return self

   def __exit__(self, exc_type, exc_value, traceback):
     self.httpd.shutdown()
     self.thread.join()

   def getURL(self):
     return 'http://{}:{}'.format(*self.httpd.server_address)


 class _Handler(http.server.SimpleHTTPRequestHandler):
   """A SimpleHTTPRequestHandler with authentication."""

   # Note: until Python 3.6, SimpleHTTPRequestHandler was only able to serve
   # files from the working directory. A 'directory' parameter was added in
   # Python 3.7, but sadly our CI builds are stuck with Python 3.6. Instead,
   # we monkey-patch translate_path() to rewrite the path.

   def __init__(self, root_directory, expected_auth, *args, **kwargs):
     self.root_directory = root_directory
     self.expected_auth = expected_auth
     super().__init__(*args, **kwargs)

   def translate_path(self, path):
     abs_path = super().translate_path(path)
     rel_path = os.path.relpath(abs_path, os.getcwd())
     return os.path.join(self.root_directory, rel_path)

   def check_auth(self):
     auth_header = self.headers.get('Authorization', None)
     if auth_header != self.expected_auth:
       self.send_error(http.HTTPStatus.UNAUTHORIZED)
       return False
     return True

   def do_HEAD(self):
     if self.check_auth():
       return super().do_HEAD()

   def do_GET(self):
     if self.check_auth():
       return super().do_GET()
	# pylint: disable=invalid-name
	# pylint: disable=g-long-ternary
	# Copyright 2021 The Bazel Authors. All rights reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.
	"""Test utils for Bzlmod."""

	import base64
	import functools
	import hashlib
	import http.server
	import json
	import os
	import pathlib
	import shutil
	import threading
	import urllib.request
	import zipfile


	def download(url):
	"""Download a file and return its content in bytes."""
	response = urllib.request.urlopen(url)
	return response.read()


	def read(path):
	"""Read a file and return its content in bytes."""
	with open(str(path), 'rb') as f:
	return f.read()


	def integrity(data):
	"""Calculate the integration value of the data with sha256."""
	hash_value = hashlib.sha256(data)
	return 'sha256-' + base64.b64encode(hash_value.digest()).decode()


	def scratchFile(path, lines=None):
	"""Creates a file at the given path with the given content."""
	with open(str(path), 'w') as f:
	if lines:
	for l in lines:
	f.write(l)
	f.write('\n')


	class Module:
	"""A class to represent information of a Bazel module."""

	def __init__(self, name, version):
	self.name = name
	self.version = version
	self.archive_url = None
	self.strip_prefix = ''
	self.module_dot_bazel = None
	self.patches = []
	self.patch_strip = 0
	self.archive_type = None

	def set_source(self, archive_url, strip_prefix=None):
	self.archive_url = archive_url
	self.strip_prefix = strip_prefix
	return self

	def set_module_dot_bazel(self, module_dot_bazel):
	self.module_dot_bazel = module_dot_bazel
	return self

	def set_patches(self, patches, patch_strip):
	self.patches = patches
	self.patch_strip = patch_strip
	return self

	def set_archive_type(self, archive_type):
	self.archive_type = archive_type
	return self


	class BazelRegistry:
	"""A class to help create a Bazel module project from scatch and add it into the registry."""

	def __init__(self, root, registry_suffix=''):
	self.root = pathlib.Path(root)
	self.projects = self.root.joinpath('projects')
	self.projects.mkdir(parents=True, exist_ok=True)
	self.archives = self.root.joinpath('archives')
	self.archives.mkdir(parents=True, exist_ok=True)
	self.registry_suffix = registry_suffix

	def setModuleBasePath(self, module_base_path):
	bazel_registry = {
	'module_base_path': module_base_path,
	}
	with self.root.joinpath('bazel_registry.json').open('w') as f:
	json.dump(bazel_registry, f, indent=4, sort_keys=True)

	def getURL(self):
	"""Return the URL of this registry."""
	return self.root.resolve().as_uri()

	def generateCcSource(
	self,
	name,
	version,
	deps=None,
	repo_names=None,
	extra_module_file_contents=None,
	):
	"""Generate a cc project with given dependency information.

	1. The cc projects implements a hello_<lib_name> function.
	2. The hello_<lib_name> function calls the same function of its
	dependencies.
	3. The hello_<lib_name> function prints "<caller name> =>
	<lib_name@version>".
	4. The BUILD file references the dependencies as their desired repo names.

	Args:
	name: The module name.
	version: The module version.
	deps: The dependencies of this module.
	repo_names: The desired repository name for some dependencies.
	extra_module_file_contents: Extra lines to append to the MODULE.bazel
	file.

	Returns:
	The generated source directory.
	"""

	src_dir = self.projects.joinpath(name, version)
	src_dir.mkdir(parents=True, exist_ok=True)
	if not deps:
	deps = {}
	if not repo_names:
	repo_names = {}
	for dep in deps:
	if dep not in repo_names:
	repo_names[dep] = dep
	if not extra_module_file_contents:
	extra_module_file_contents = []

	def calc_repo_name_str(dep):
	if dep == repo_names[dep]:
	return ''
	return ', repo_name = "%s"' % repo_names[dep]

	scratchFile(src_dir.joinpath('WORKSPACE'))
	scratchFile(
	src_dir.joinpath('MODULE.bazel'),
	[
	'module(',
	' name = "%s",' % name,
	' version = "%s",' % version,
	' compatibility_level = 1,',
	')',
	]
	+ [
	'bazel_dep(name = "%s", version = "%s"%s)'
	% (dep, version, calc_repo_name_str(dep))
	for dep, version in deps.items()
	]
	+ extra_module_file_contents,
	)

	scratchFile(
	src_dir.joinpath(name.lower() + '.h'), [
	'#ifndef %s_H' % name.upper(),
	'#define %s_H' % name.upper(),
	'#include <string>',
	'void hello_%s(const std::string& caller);' % name.lower(),
	'#endif',
	])
	scratchFile(
	src_dir.joinpath(name.lower() + '.cc'), [
	'#include <stdio.h>',
	'#include "%s.h"' % name.lower(),
	] + ['#include "%s.h"' % dep.lower() for dep in deps] + [
	'void hello_%s(const std::string& caller) {' % name.lower(),
	' std::string lib_name = "%s@%s%s";' %
	(name, version, self.registry_suffix),
	' printf("%s => %s\\n", caller.c_str(), lib_name.c_str());',
	] + [' hello_%s(lib_name);' % dep.lower() for dep in deps] + [
	'}',
	])
	scratchFile(
	src_dir.joinpath('BUILD'), [
	'package(default_visibility = ["//visibility:public"])',
	'cc_library(',
	' name = "lib_%s",' % name.lower(),
	' srcs = ["%s.cc"],' % name.lower(),
	' hdrs = ["%s.h"],' % name.lower(),
	] + ([
	' deps = ["%s"],' % ('", "'.join([
	'@%s//:lib_%s' % (repo_names[dep], dep.lower()) for dep in deps
	])),
	] if deps else []) + [
	')',
	])
	return src_dir

	def createArchive(self, name, version, src_dir, filename_pattern='%s.%s.zip'):
	"""Create an archive with a given source directory."""
	zip_path = self.archives.joinpath(filename_pattern % (name, version))
	zip_obj = zipfile.ZipFile(str(zip_path), 'w')
	for foldername, _, filenames in os.walk(str(src_dir)):
	for filename in filenames:
	filepath = os.path.join(foldername, filename)
	zip_obj.write(filepath,
	str(pathlib.Path(filepath).relative_to(src_dir)))
	zip_obj.close()
	return zip_path

	def addModule(self, module):
	"""Add a module into the registry."""
	module_dir = self.root.joinpath('modules', module.name, module.version)
	module_dir.mkdir(parents=True, exist_ok=True)

	# Copy MODULE.bazel to the registry
	module_dot_bazel = module_dir.joinpath('MODULE.bazel')
	shutil.copy(str(module.module_dot_bazel), str(module_dot_bazel))

	# Create source.json & copy patch files to the registry
	source = {
	'url': module.archive_url,
	'integrity': integrity(download(module.archive_url)),
	}
	if module.strip_prefix:
	source['strip_prefix'] = module.strip_prefix

	if module.patches:
	patch_dir = module_dir.joinpath('patches')
	patch_dir.mkdir()
	source['patches'] = {}
	source['patch_strip'] = module.patch_strip
	for patch_path in module.patches:
	patch = pathlib.Path(patch_path)
	source['patches'][patch.name] = integrity(read(patch))
	shutil.copy(str(patch), str(patch_dir))

	if module.archive_type:
	source['archive_type'] = module.archive_type

	with module_dir.joinpath('source.json').open('w') as f:
	json.dump(source, f, indent=4, sort_keys=True)

	def createCcModule(
	self,
	name,
	version,
	deps=None,
	repo_names=None,
	patches=None,
	patch_strip=0,
	archive_pattern=None,
	archive_type=None,
	extra_module_file_contents=None,
	):
	"""Generate a cc project and add it as a module into the registry."""
	src_dir = self.generateCcSource(
	name, version, deps, repo_names, extra_module_file_contents
	)
	if archive_pattern:
	archive = self.createArchive(
	name, version, src_dir, filename_pattern=archive_pattern
	)
	else:
	archive = self.createArchive(name, version, src_dir)
	module = Module(name, version)
	module.set_source(archive.resolve().as_uri())
	module.set_module_dot_bazel(src_dir.joinpath('MODULE.bazel'))
	if patches:
	module.set_patches(patches, patch_strip)
	if archive_type:
	module.set_archive_type(archive_type)

	self.addModule(module)
	return self

	def addMetadata(self,
	name,
	homepage=None,
	maintainers=None,
	versions=None,
	yanked_versions=None):
	"""Generate a module metadata file and add it to the registry."""
	if maintainers is None:
	maintainers = []
	if versions is None:
	versions = []
	if yanked_versions is None:
	yanked_versions = {}

	module_dir = self.root.joinpath('modules', name)
	module_dir.mkdir(parents=True, exist_ok=True)

	metadata = {
	'homepage': homepage,
	'maintainers': maintainers,
	'versions': versions,
	'yanked_versions': yanked_versions
	}

	with module_dir.joinpath('metadata.json').open('w') as f:
	json.dump(metadata, f, indent=4, sort_keys=True)

	return self

	def createLocalPathModule(self, name, version, path, deps=None):
	"""Add a local module into the registry."""
	module_dir = self.root.joinpath('modules', name, version)
	module_dir.mkdir(parents=True, exist_ok=True)

	# Create source.json & copy patch files to the registry
	source = {
	'type': 'local_path',
	'path': path,
	}

	if deps is None:
	deps = {}

	module_file_lines = [
	'module(',
	' name = "%s",' % name,
	' version = "%s",' % version,
	')',
	] + ['bazel_dep(name="%s",version="%s")' % p for p in deps.items()]
	scratchFile(module_dir.joinpath('MODULE.bazel'), module_file_lines)
	self.projects.joinpath(path).mkdir(exist_ok=True)
	scratchFile(self.projects.joinpath(path, 'MODULE.bazel'), module_file_lines)

	with module_dir.joinpath('source.json').open('w') as f:
	json.dump(source, f, indent=4, sort_keys=True)


	class StaticHTTPServer:
	"""An HTTP server serving static files, optionally with authentication."""

	def __init__(self, root_directory, expected_auth=None):
	self.root_directory = root_directory
	self.expected_auth = expected_auth

	def __enter__(self):
	address = ('localhost', 0) # assign random port
	handler = functools.partial(
	_Handler, self.root_directory, self.expected_auth
	)
	self.httpd = http.server.HTTPServer(address, handler)
	self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True)
	self.thread.start()
	return self

	def __exit__(self, exc_type, exc_value, traceback):
	self.httpd.shutdown()
	self.thread.join()

	def getURL(self):
	return 'http://{}:{}'.format(*self.httpd.server_address)


	class _Handler(http.server.SimpleHTTPRequestHandler):
	"""A SimpleHTTPRequestHandler with authentication."""

	# Note: until Python 3.6, SimpleHTTPRequestHandler was only able to serve
	# files from the working directory. A 'directory' parameter was added in
	# Python 3.7, but sadly our CI builds are stuck with Python 3.6. Instead,
	# we monkey-patch translate_path() to rewrite the path.

	def __init__(self, root_directory, expected_auth, args, *kwargs):
	self.root_directory = root_directory
	self.expected_auth = expected_auth
	super().__init__(args, *kwargs)

	def translate_path(self, path):
	abs_path = super().translate_path(path)
	rel_path = os.path.relpath(abs_path, os.getcwd())
	return os.path.join(self.root_directory, rel_path)

	def check_auth(self):
	auth_header = self.headers.get('Authorization', None)
	if auth_header != self.expected_auth:
	self.send_error(http.HTTPStatus.UNAUTHORIZED)
	return False
	return True

	def do_HEAD(self):
	if self.check_auth():
	return super().do_HEAD()

	def do_GET(self):
	if self.check_auth():
	return super().do_GET()