Skip to content

chore: refactor training job tests to configure temp test directory #1254

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sararob merged 5 commits into from
May 26, 2022
Merged

chore: refactor training job tests to configure temp test directory #1254

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8e9a20e
replace filepath in test with tmpdir var
sararob May 20, 2022
93d3261
update script path var in training job tests
sararob May 24, 2022
9cd150c
Merge branch 'main' into sr-unit-test-updates
sararob May 24, 2022
95746ff
run linter
sararob May 25, 2022
680b133
Merge branch 'main' into sr-unit-test-updates
sararob May 26, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 13 additions & 9 deletions tests/unit/aiplatform/test_training_jobs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

from distutils import core
import copy
import os
import functools
import importlib
import logging
Expand Down Expand Up @@ -73,7 +74,8 @@
_TEST_GCS_PATH = f"{_TEST_BUCKET_NAME}/{_TEST_GCS_PATH_WITHOUT_BUCKET}"
_TEST_GCS_PATH_WITH_TRAILING_SLASH = f"{_TEST_GCS_PATH}/"
_TEST_LOCAL_SCRIPT_FILE_NAME = "____test____script.py"
_TEST_LOCAL_SCRIPT_FILE_PATH = f"path/to/{_TEST_LOCAL_SCRIPT_FILE_NAME}"
_TEST_TEMPDIR = tempfile.mkdtemp()
_TEST_LOCAL_SCRIPT_FILE_PATH = os.path.join(_TEST_TEMPDIR, _TEST_LOCAL_SCRIPT_FILE_NAME)
Comment on lines +77 to +78
Copy link
Contributor

@ucdmkt ucdmkt May 25, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you think you can afford dependency on absl-py? If so, standardizing on tmp dir handling, with ability to inject via a public envvar, is beneficial.

https://github.com/abseil/abseil-py/blob/main/absl/testing/absltest.py#L158-L165

Alternatively, pytest also has a built in fixture to do so in a standard way , but I understand that pytest is used only for test driver, but not as a unittest utility framework as it is now, so it will be anyways additional dependency. I misspoke on this one. pytest is also used as utility.

https://docs.pytest.org/en/6.2.x/tmpdir.html#base-temporary-directory

_TEST_PYTHON_SOURCE = """
print('hello world')
"""
Expand Down Expand Up @@ -449,11 +451,11 @@ class TestTrainingScriptPythonPackager:
def setup_method(self):
importlib.reload(initializer)
importlib.reload(aiplatform)
with open(_TEST_LOCAL_SCRIPT_FILE_NAME, "w") as fp:
with open(_TEST_LOCAL_SCRIPT_FILE_PATH, "w") as fp:
fp.write(_TEST_PYTHON_SOURCE)

def teardown_method(self):
pathlib.Path(_TEST_LOCAL_SCRIPT_FILE_NAME).unlink()
pathlib.Path(_TEST_LOCAL_SCRIPT_FILE_PATH).unlink()
python_package_file = f"{source_utils._TrainingScriptPythonPackager._ROOT_MODULE}-{source_utils._TrainingScriptPythonPackager._SETUP_PY_VERSION}.tar.gz"
if pathlib.Path(python_package_file).is_file():
pathlib.Path(python_package_file).unlink()
Expand All @@ -467,14 +469,14 @@ def teardown_method(self):
)

def test_packager_creates_and_copies_python_package(self):
tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_NAME)
tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_PATH)
tsp.package_and_copy(copy_method=local_copy_method)
assert pathlib.Path(
f"{tsp._ROOT_MODULE}-{tsp._SETUP_PY_VERSION}.tar.gz"
).is_file()

def test_created_package_module_is_installable_and_can_be_run(self):
tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_NAME)
tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_PATH)
source_dist_path = tsp.package_and_copy(copy_method=local_copy_method)
subprocess.check_output(["pip3", "install", source_dist_path])
module_output = subprocess.check_output(
Expand All @@ -484,7 +486,7 @@ def test_created_package_module_is_installable_and_can_be_run(self):

def test_requirements_are_in_package(self):
tsp = source_utils._TrainingScriptPythonPackager(
_TEST_LOCAL_SCRIPT_FILE_NAME, requirements=_TEST_REQUIREMENTS
_TEST_LOCAL_SCRIPT_FILE_PATH, requirements=_TEST_REQUIREMENTS
)
source_dist_path = tsp.package_and_copy(copy_method=local_copy_method)
with tarfile.open(source_dist_path) as tf:
Expand All @@ -503,15 +505,15 @@ def test_packaging_fails_whith_RuntimeError(self):
mock_subprocess.returncode = 1
mock_popen.return_value = mock_subprocess
tsp = source_utils._TrainingScriptPythonPackager(
_TEST_LOCAL_SCRIPT_FILE_NAME
_TEST_LOCAL_SCRIPT_FILE_PATH
)
with pytest.raises(RuntimeError):
tsp.package_and_copy(copy_method=local_copy_method)

def test_package_and_copy_to_gcs_copies_to_gcs(self, mock_client_bucket):
mock_client_bucket, mock_blob = mock_client_bucket

tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_NAME)
tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_PATH)

gcs_path = tsp.package_and_copy_to_gcs(
gcs_staging_dir=_TEST_BUCKET_NAME, project=_TEST_PROJECT
Expand Down Expand Up @@ -838,7 +840,9 @@ class TestCustomTrainingJob:
def setup_method(self):
importlib.reload(initializer)
importlib.reload(aiplatform)
self._local_script_file_name = f"{uuid.uuid4()}-{_TEST_LOCAL_SCRIPT_FILE_NAME}"
self._local_script_file_name = os.path.join(
_TEST_TEMPDIR, f"{uuid.uuid4()}-{_TEST_LOCAL_SCRIPT_FILE_NAME}"
)
with open(self._local_script_file_name, "w") as fp:
fp.write(_TEST_PYTHON_SOURCE)

Expand Down