From 5b753eaa9c3e8be703274cd398ecb7a47ae151c3 Mon Sep 17 00:00:00 2001
From: Asaf Gardin <asafg@ai21.com>
Date: Tue, 12 Dec 2023 15:15:24 +0200
Subject: [PATCH] test: sagemaker tests

---
 .../sagemaker/{answer_sm.py => answer.py}     |  0
 .../{completion_sm.py => completion.py}       |  0
 examples/sagemaker/{gec_sm.py => gec.py}      |  0
 .../{paraphrase_sm.py => paraphrase.py}       |  0
 .../{summarization_sm.py => summarization.py} |  0
 scripts/run_sm_examples.py                    | 23 -------------
 tests/integration_tests/test_sagemaker.py     | 33 +++++++++++++++++++
 7 files changed, 33 insertions(+), 23 deletions(-)
 rename examples/sagemaker/{answer_sm.py => answer.py} (100%)
 rename examples/sagemaker/{completion_sm.py => completion.py} (100%)
 rename examples/sagemaker/{gec_sm.py => gec.py} (100%)
 rename examples/sagemaker/{paraphrase_sm.py => paraphrase.py} (100%)
 rename examples/sagemaker/{summarization_sm.py => summarization.py} (100%)
 delete mode 100644 scripts/run_sm_examples.py
 create mode 100644 tests/integration_tests/test_sagemaker.py

diff --git a/examples/sagemaker/answer_sm.py b/examples/sagemaker/answer.py
similarity index 100%
rename from examples/sagemaker/answer_sm.py
rename to examples/sagemaker/answer.py
diff --git a/examples/sagemaker/completion_sm.py b/examples/sagemaker/completion.py
similarity index 100%
rename from examples/sagemaker/completion_sm.py
rename to examples/sagemaker/completion.py
diff --git a/examples/sagemaker/gec_sm.py b/examples/sagemaker/gec.py
similarity index 100%
rename from examples/sagemaker/gec_sm.py
rename to examples/sagemaker/gec.py
diff --git a/examples/sagemaker/paraphrase_sm.py b/examples/sagemaker/paraphrase.py
similarity index 100%
rename from examples/sagemaker/paraphrase_sm.py
rename to examples/sagemaker/paraphrase.py
diff --git a/examples/sagemaker/summarization_sm.py b/examples/sagemaker/summarization.py
similarity index 100%
rename from examples/sagemaker/summarization_sm.py
rename to examples/sagemaker/summarization.py
diff --git a/scripts/run_sm_examples.py b/scripts/run_sm_examples.py
deleted file mode 100644
index 4236dd8a..00000000
--- a/scripts/run_sm_examples.py
+++ /dev/null
@@ -1,23 +0,0 @@
-"""
-You need to be authenticated to the AWS when you run this and make sure the endpoints are there
-"""
-
-import subprocess
-from pathlib import Path
-
-SAGEMAKER_EXAMPLE_FILES = [
-    "answer_sm.py",
-    "completion_sm.py",
-    "completion_sm_destination.py",
-    "completion_sm_injected_session.py",
-    "gec_sm.py",
-    # 'launch_sm_model.py',
-    "paraphrase_sm.py",
-    "summarization_sm.py",
-]
-
-SAGEMAKER_DIR = "sagemaker"
-
-for example_file_name in SAGEMAKER_EXAMPLE_FILES:
-    if subprocess.call(["python", Path("..") / "examples" / SAGEMAKER_DIR / example_file_name]):
-        raise Exception(f"failed to run {example_file_name}")
diff --git a/tests/integration_tests/test_sagemaker.py b/tests/integration_tests/test_sagemaker.py
new file mode 100644
index 00000000..21bc087f
--- /dev/null
+++ b/tests/integration_tests/test_sagemaker.py
@@ -0,0 +1,33 @@
+import subprocess
+
+import pytest
+from pathlib import Path
+
+SAGEMAKER_DIR = "sagemaker"
+
+SAGEMAKER_PATH = Path(__file__).parent.parent.parent / "examples" / SAGEMAKER_DIR
+
+
+@pytest.mark.skip(reason="SageMaker integration tests need endpoints to be running")
+@pytest.mark.parametrize(
+    argnames=["test_file_name"],
+    argvalues=[
+        ("answer.py",),
+        ("completion.py",),
+        ("gec.py",),
+        ("paraphrase.py",),
+        ("summarization.py",),
+    ],
+    ids=[
+        "when_answer__should_return_ok",
+        "when_completion__should_return_ok",
+        "when_gec__should_return_ok",
+        "when_paraphrase__should_return_ok",
+        "when_summarization__should_return_ok",
+    ],
+)
+def test_sagemaker(test_file_name: str):
+    file_path = SAGEMAKER_PATH / test_file_name
+    print(f"About to run: {file_path}")
+    exit_code = subprocess.call(["python", file_path])
+    assert exit_code == 0, f"failed to run {test_file_name}"