From 6f69b19ff583a37387b8403f15fae7bfbcede4ba Mon Sep 17 00:00:00 2001
From: Bagatur <baskaryan@gmail.com>
Date: Thu, 17 Aug 2023 16:45:52 -0700
Subject: [PATCH] wip tests

---
 .../unit_tests/schema/runnable/__init__.py    |  0
 .../unit_tests/schema/runnable/test_locals.py | 31 +++++++++++++++++++
 .../schema/{ => runnable}/test_runnable.py    |  0
 3 files changed, 31 insertions(+)
 create mode 100644 libs/langchain/tests/unit_tests/schema/runnable/__init__.py
 create mode 100644 libs/langchain/tests/unit_tests/schema/runnable/test_locals.py
 rename libs/langchain/tests/unit_tests/schema/{ => runnable}/test_runnable.py (100%)

diff --git a/libs/langchain/tests/unit_tests/schema/runnable/__init__.py b/libs/langchain/tests/unit_tests/schema/runnable/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/libs/langchain/tests/unit_tests/schema/runnable/test_locals.py b/libs/langchain/tests/unit_tests/schema/runnable/test_locals.py
new file mode 100644
index 00000000000..d0a3fb38d9d
--- /dev/null
+++ b/libs/langchain/tests/unit_tests/schema/runnable/test_locals.py
@@ -0,0 +1,31 @@
+import pytest
+
+from langchain.schema.runnable import GetLocalVar, PutLocalVar
+
+
+@pytest.mark.asyncio
+async def test_put_get() -> None:
+    runnable = PutLocalVar("input") | GetLocalVar("input")
+    assert runnable.invoke("foo") == "foo"
+    assert runnable.batch(["foo", "bar"]) == ["foo", "bar"]
+    assert list(runnable.stream("foo"))[0] == "foo"
+
+    assert await runnable.ainvoke("foo") == "foo"
+    assert await runnable.abatch(["foo", "bar"]) == ["foo", "bar"]
+    async for x in runnable.astream("foo"):
+        assert x == "foo"
+
+
+def test_missing_config() -> None:
+    with pytest.raises(ValueError):
+        PutLocalVar("input").invoke("foo")
+
+    with pytest.raises(ValueError):
+        GetLocalVar("input").invoke("foo")
+
+
+def test_get_missing_var_invoke() -> None:
+    runnable = PutLocalVar("input") | GetLocalVar("missing")
+    with pytest.raises(KeyError):
+        runnable.invoke("foo")
+
diff --git a/libs/langchain/tests/unit_tests/schema/test_runnable.py b/libs/langchain/tests/unit_tests/schema/runnable/test_runnable.py
similarity index 100%
rename from libs/langchain/tests/unit_tests/schema/test_runnable.py
rename to libs/langchain/tests/unit_tests/schema/runnable/test_runnable.py