From 8f58ed3d70fefdf2c8442137c91de16287344514 Mon Sep 17 00:00:00 2001
From: Piyush Jain <piyushjain@duck.com>
Date: Thu, 28 Mar 2024 15:33:59 -0700
Subject: [PATCH] Added integration tests.

---
 libs/aws/pyproject.toml                       |  3 +++
 .../integration_tests/chat_models/__init__.py |  0
 .../integration_tests/graphs/__init__.py      |  0
 .../tests/integration_tests/llms/__init__.py  |  0
 .../llms/test_sagemaker_endpoint.py           | 21 +++++++++++++++++++
 .../tests/integration_tests/test_compile.py   |  7 +++++++
 6 files changed, 31 insertions(+)
 create mode 100644 libs/aws/tests/integration_tests/chat_models/__init__.py
 create mode 100644 libs/aws/tests/integration_tests/graphs/__init__.py
 create mode 100644 libs/aws/tests/integration_tests/llms/__init__.py
 create mode 100644 libs/aws/tests/integration_tests/llms/test_sagemaker_endpoint.py
 create mode 100644 libs/aws/tests/integration_tests/test_compile.py

diff --git a/libs/aws/pyproject.toml b/libs/aws/pyproject.toml
index 3d60dab8..33588913 100644
--- a/libs/aws/pyproject.toml
+++ b/libs/aws/pyproject.toml
@@ -84,6 +84,9 @@ addopts = "--snapshot-warn-unused --strict-markers --strict-config --durations=5
 # Registering custom markers.
 # https://docs.pytest.org/en/7.1.x/example/markers.html#registering-markers
 markers = [
+  "requires: mark tests as requiring a specific library",
+  "asyncio: mark tests as requiring asyncio",
   "compile: mark placeholder test used to compile integration tests without running them",
+  "scheduled: mark tests to run in scheduled testing",
 ]
 asyncio_mode = "auto"
diff --git a/libs/aws/tests/integration_tests/chat_models/__init__.py b/libs/aws/tests/integration_tests/chat_models/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/libs/aws/tests/integration_tests/graphs/__init__.py b/libs/aws/tests/integration_tests/graphs/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/libs/aws/tests/integration_tests/llms/__init__.py b/libs/aws/tests/integration_tests/llms/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/libs/aws/tests/integration_tests/llms/test_sagemaker_endpoint.py b/libs/aws/tests/integration_tests/llms/test_sagemaker_endpoint.py
new file mode 100644
index 00000000..cb5e32ef
--- /dev/null
+++ b/libs/aws/tests/integration_tests/llms/test_sagemaker_endpoint.py
@@ -0,0 +1,21 @@
+from typing import Dict
+
+from langchain_aws.llms import SagemakerEndpoint
+from langchain_aws.llms.sagemaker_endpoint import LLMContentHandler
+
+
+class ContentHandler(LLMContentHandler):
+    def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:
+        return b""
+
+    def transform_output(self, output: bytes) -> str:
+        return ""
+
+
+def test_sagemaker_endpoint_name_param() -> None:
+    llm = SagemakerEndpoint(
+        endpoint_name="foo",
+        content_handler=ContentHandler(),
+        region_name="us-west-2",
+    )
+    assert llm.endpoint_name == "foo"
diff --git a/libs/aws/tests/integration_tests/test_compile.py b/libs/aws/tests/integration_tests/test_compile.py
new file mode 100644
index 00000000..33ecccdf
--- /dev/null
+++ b/libs/aws/tests/integration_tests/test_compile.py
@@ -0,0 +1,7 @@
+import pytest
+
+
+@pytest.mark.compile
+def test_placeholder() -> None:
+    """Used for compiling integration tests without running any real tests."""
+    pass