test: sagemaker tests

AI21Labs · Dec 12, 2023 · 5b753ea · 5b753ea
1 parent 71d1cc9
commit 5b753ea
Show file tree

Hide file tree

Showing 7 changed files with 33 additions and 23 deletions.
diff --git a/examples/sagemaker/answer_sm.py → examples/sagemaker/answer.py b/examples/sagemaker/answer_sm.py → examples/sagemaker/answer.py
diff --git a/examples/sagemaker/completion_sm.py → examples/sagemaker/completion.py b/examples/sagemaker/completion_sm.py → examples/sagemaker/completion.py
diff --git a/examples/sagemaker/gec_sm.py → examples/sagemaker/gec.py b/examples/sagemaker/gec_sm.py → examples/sagemaker/gec.py
diff --git a/examples/sagemaker/paraphrase_sm.py → examples/sagemaker/paraphrase.py b/examples/sagemaker/paraphrase_sm.py → examples/sagemaker/paraphrase.py
diff --git a/examples/sagemaker/summarization_sm.py → examples/sagemaker/summarization.py b/examples/sagemaker/summarization_sm.py → examples/sagemaker/summarization.py
diff --git a/scripts/run_sm_examples.py b/scripts/run_sm_examples.py
diff --git a/tests/integration_tests/test_sagemaker.py 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}"