136 lines
5.8 KiB
Python
136 lines
5.8 KiB
Python
# Copyright 2019 Google LLC
|
|
#
|
|
# 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 speci
|
|
|
|
import os
|
|
import re
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
import mock
|
|
|
|
from kfp.containers import build_image_from_working_dir
|
|
|
|
|
|
class MockImageBuilder:
|
|
def __init__(self, dockerfile_text_check : Callable[[str], None] = None, requirements_text_check : Callable[[str], None] = None, file_paths_check : Callable[[str], None] = None):
|
|
self.dockerfile_text_check = dockerfile_text_check
|
|
self.requirements_text_check = requirements_text_check
|
|
self.file_paths_check = file_paths_check
|
|
|
|
def build(self, local_dir = None, target_image = None, timeout = 1000):
|
|
if self.dockerfile_text_check:
|
|
actual_dockerfile_text = (Path(local_dir) / 'Dockerfile').read_text()
|
|
self.dockerfile_text_check(actual_dockerfile_text)
|
|
if self.requirements_text_check:
|
|
actual_requirements_text = (Path(local_dir) / 'requirements.txt').read_text()
|
|
self.requirements_text_check(actual_requirements_text)
|
|
if self.file_paths_check:
|
|
file_paths = set(os.path.relpath(os.path.join(dirpath, file_name), local_dir) for dirpath, dirnames, filenames in os.walk(local_dir) for file_name in filenames)
|
|
self.file_paths_check(file_paths)
|
|
return target_image
|
|
|
|
|
|
class BuildImageApiTests(unittest.TestCase):
|
|
def test_build_image_from_working_dir(self):
|
|
expected_dockerfile_text_re = '''
|
|
FROM python:3.6.5
|
|
WORKDIR /.*
|
|
COPY requirements.txt .
|
|
RUN python3 -m pip install -r requirements.txt
|
|
COPY . .
|
|
'''
|
|
#mock_builder =
|
|
with tempfile.TemporaryDirectory() as context_dir:
|
|
requirements_text = 'pandas==1.24'
|
|
requirements_txt_relpath = Path('.') / 'requirements.txt'
|
|
file1_py_relpath = Path('.') / 'lib' / 'file1.py'
|
|
file1_sh_relpath = Path('.') / 'lib' / 'file1.sh'
|
|
|
|
context_path = Path(context_dir)
|
|
(context_path / requirements_txt_relpath).write_text(requirements_text)
|
|
(context_path / file1_py_relpath).parent.mkdir(parents=True, exist_ok=True)
|
|
(context_path / file1_py_relpath).write_text('#py file')
|
|
(context_path / file1_sh_relpath).parent.mkdir(parents=True, exist_ok=True)
|
|
(context_path / file1_sh_relpath).write_text('#sh file')
|
|
expected_file_paths = {
|
|
'Dockerfile',
|
|
str(requirements_txt_relpath),
|
|
str(file1_py_relpath),
|
|
}
|
|
def dockerfile_text_check(actual_dockerfile_text):
|
|
self.assertRegex(actual_dockerfile_text.strip(), expected_dockerfile_text_re.strip())
|
|
def requirements_text_check(actual_requirements_text):
|
|
self.assertEqual(actual_requirements_text.strip(), requirements_text.strip())
|
|
def file_paths_check(file_paths):
|
|
self.assertEqual(file_paths, expected_file_paths)
|
|
|
|
builder = MockImageBuilder(dockerfile_text_check, requirements_text_check, file_paths_check)
|
|
result = build_image_from_working_dir(working_dir=context_dir, base_image='python:3.6.5', builder=builder)
|
|
|
|
|
|
def test_image_cache(self):
|
|
builder = InvocationCountingDummyImageBuilder()
|
|
|
|
from kfp.containers._cache import clear_cache
|
|
clear_cache('build_image_from_working_dir')
|
|
|
|
self.assertEqual(builder.invocations_count, 0)
|
|
with prepare_context_dir(py_content='py1', sh_content='sh1') as context_dir:
|
|
build_image_from_working_dir(working_dir=context_dir, base_image='python:3.6.5', builder=builder)
|
|
self.assertEqual(builder.invocations_count, 1)
|
|
|
|
# Check that changes to .sh files do not break cache
|
|
with prepare_context_dir(py_content='py1', sh_content='sh2') as context_dir:
|
|
build_image_from_working_dir(working_dir=context_dir, base_image='python:3.6.5', builder=builder)
|
|
self.assertEqual(builder.invocations_count, 1)
|
|
|
|
# Check that changes to .py files result in new image being built
|
|
with prepare_context_dir(py_content='py2', sh_content='sh1') as context_dir:
|
|
build_image_from_working_dir(working_dir=context_dir, base_image='python:3.6.5', builder=builder)
|
|
self.assertEqual(builder.invocations_count, 2)
|
|
|
|
|
|
class InvocationCountingDummyImageBuilder:
|
|
def __init__(self):
|
|
self.invocations_count = 0
|
|
|
|
def build(self, local_dir = None, target_image = None, timeout = 1000):
|
|
self.invocations_count = self.invocations_count + 1
|
|
return "image/name@sha256:0123456789abcdef0123456789abcdef"
|
|
|
|
|
|
def prepare_context_dir(py_content: str = '#py file', sh_content: str = '#sh file') -> str:
|
|
context_dir = tempfile.TemporaryDirectory()
|
|
#Preparing context
|
|
requirements_text = 'pandas==1.24'
|
|
requirements_txt_relpath = Path('.') / 'requirements.txt'
|
|
file1_py_relpath = Path('.') / 'lib' / 'file1.py'
|
|
file1_sh_relpath = Path('.') / 'lib' / 'file1.sh'
|
|
|
|
context_path = Path(context_dir.name)
|
|
(context_path / requirements_txt_relpath).write_text(requirements_text)
|
|
(context_path / file1_py_relpath).parent.mkdir(parents=True, exist_ok=True)
|
|
(context_path / file1_py_relpath).write_text(py_content)
|
|
(context_path / file1_sh_relpath).parent.mkdir(parents=True, exist_ok=True)
|
|
(context_path / file1_sh_relpath).write_text(sh_content)
|
|
# End preparing context
|
|
|
|
return context_dir
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|