Skip to content

Commit

Permalink
chore: refactor training job tests to configure temp test directory (#…
Browse files Browse the repository at this point in the history
…1254)

* replace filepath in test with tmpdir var

* update script path var in training job tests

* run linter
  • Loading branch information
sararob authored May 26, 2022
1 parent 0ecfe1e commit 74dbabd
Showing 1 changed file with 13 additions and 9 deletions.
22 changes: 13 additions & 9 deletions tests/unit/aiplatform/test_training_jobs.py
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)
_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

0 comments on commit 74dbabd

Please sign in to comment.