fix test

Author: nadiaya

test-toolkit/integration/__init__.py

@@ -19,12 +19,19 @@
 mnist_script = os.path.join(mnist_path, 'mnist.py')
 data_dir = os.path.join(mnist_path, 'data')
 training_dir = os.path.join(data_dir, 'training')
+cpu_sub_dir = 'model_cpu'
+gpu_sub_dir = 'model_gpu'
+eia_sub_dir = 'model_eia'
 
-model_cpu_dir = os.path.join(mnist_path, 'model_cpu')
+model_cpu_dir = os.path.join(mnist_path, cpu_sub_dir)
+mnist_cpu_script = os.path.join(model_cpu_dir, 'mnist.py')
 model_cpu_1d_dir = os.path.join(model_cpu_dir, '1d')
 mnist_1d_script = os.path.join(model_cpu_1d_dir, 'mnist_1d.py')
-model_gpu_dir = os.path.join(mnist_path, 'model_gpu')
+model_gpu_dir = os.path.join(mnist_path, gpu_sub_dir)
+mnist_gpu_script = os.path.join(model_gpu_dir, 'mnist.py')
 model_gpu_1d_dir = os.path.join(model_gpu_dir, '1d')
+model_eia_dir = os.path.join(mnist_path, eia_sub_dir)
+mnist_eia_script = os.path.join(model_eia_dir, 'mnist.py')
 call_model_fn_once_script = os.path.join(resources_path, 'call_model_fn_once.py')
 
 ROLE = 'dummy/unused-role'

test-toolkit/integration/sagemaker/test_mnist.py

@@ -19,8 +19,8 @@
 import sagemaker
 from sagemaker.pytorch import PyTorchModel
 
-from test.integration import model_cpu_dir, mnist_cpu_script, mnist_gpu_script, model_eia_dir, mnist_eia_script
-from test.integration.sagemaker.timeout import timeout_and_delete_endpoint
+from integration import model_cpu_dir, mnist_cpu_script, mnist_gpu_script, model_eia_dir, mnist_eia_script
+from integration.sagemaker.timeout import timeout_and_delete_endpoint
 
 
 @pytest.mark.cpu_test