From 87b5a84cfb86d3e2ad059c81f70f7791a275f4b4 Mon Sep 17 00:00:00 2001
From: blob42 <spike@w530>
Date: Thu, 2 Mar 2023 19:33:48 +0100
Subject: [PATCH] update tests and docstrings

---
 langchain/utilities/docker/__init__.py | 33 +++++++++++++-------------
 tests/unit_tests/test_docker.py        |  6 ++---
 2 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/langchain/utilities/docker/__init__.py b/langchain/utilities/docker/__init__.py
index 4020b276..92c34703 100644
--- a/langchain/utilities/docker/__init__.py
+++ b/langchain/utilities/docker/__init__.py
@@ -17,7 +17,7 @@ import logging
 
 from .images import BaseImage, get_image_template, Python, Shell
 
-from typing import Any, Dict, Optional, Union, Type
+from typing import Any, Dict, Optional, Union, Type, List
 from pydantic import BaseModel, PrivateAttr, Extra, root_validator, validator, Field
 
 logger = logging.getLogger(__name__)
@@ -179,7 +179,7 @@ class DockerWrapper(BaseModel, extra=Extra.allow):
     Args:
         image (str | Type[BaseImage]): Docker image to use for execution. The
         image can be a string or a subclass of images.BaseImage.
-
+        default_command (List[str]): Default command to use when creating the container.
     """
 
     _docker_client: DockerClient = PrivateAttr()
@@ -187,6 +187,7 @@ class DockerWrapper(BaseModel, extra=Extra.allow):
     image: Union[str, Type[BaseImage]] = Field(default_factory=Shell,skip=True)
     from_env: Optional[bool] = Field(default=True, skip=True)
 
+
     # @property
     # def image_name(self) -> str:
     #     """The image name that will be used when creating a container."""
@@ -303,12 +304,12 @@ class DockerWrapper(BaseModel, extra=Extra.allow):
 
         # TODO: handle docker APIError ?
         except APIError as e:
-            print(f"APIError: {e}")
+            logger.debug(f"APIError: {e}")
             return "ERROR"
 
 
 
-    def exec_run(self, query: str, socket_timeout: int = 5,
+    def exec_run(self, query: str, timeout: int = 5,
                  delay: float = 0.5,
                  with_stderr: bool = False,
                  **kwargs: Any) -> str:
@@ -319,7 +320,7 @@ class DockerWrapper(BaseModel, extra=Extra.allow):
         using Docker API. It effectively simulates a tty session.
 
         Args:
-            socket_timeout (int): The timeout for receiving from the attached stdin.
+            timeout (int): The timeout for receiving from the attached stdin.
             delay (int): The delay in seconds before running the command.
             **kwargs: Pass extra parameters to DockerClient.container.exec_run.
         """
@@ -363,20 +364,20 @@ class DockerWrapper(BaseModel, extra=Extra.allow):
 
 
         # input()
-        with DockerSocket(_socket, timeout=socket_timeout) as _socket:
+        with DockerSocket(_socket, timeout=timeout) as _socket:
             # flush the output buffer (if any prompt)
-            flush = _socket.recv()
-            _socket.setblocking(True)
-            logger.debug(f"flushed output: {flush}")
-            # TEST: make sure the container is ready ? use a blocking first call
-            _socket.sendall(query.encode('utf-8'))
-
-            #NOTE: delay ensures that the command is executed after the input is sent
-            sleep(delay) #this should be available as a parameter
-
-            # read the output
             output = None
             try:
+                flush = _socket.recv()
+                _socket.setblocking(True)
+                logger.debug(f"flushed output: {flush}")
+                # TEST: make sure the container is ready ? use a blocking first call
+                _socket.sendall(query.encode('utf-8'))
+
+                #NOTE: delay ensures that the command is executed after the input is sent
+                sleep(delay) #this should be available as a parameter
+
+                # read the output
                 output = _socket.recv()
             except socket.timeout:
                 return "ERROR: timeout"
diff --git a/tests/unit_tests/test_docker.py b/tests/unit_tests/test_docker.py
index 11fb471f..0f582344 100644
--- a/tests/unit_tests/test_docker.py
+++ b/tests/unit_tests/test_docker.py
@@ -52,7 +52,7 @@ class TestDockerUtility:
         """Test inner command with non zero exit"""
         docker = DockerWrapper()
         output = docker.run('todo handle APIError')
-        assert output == 'ERROR'
+        assert str(output).startswith("STDERR") or str(output).startswith("ERROR")
 
     def test_check_gvisor_runtime(self) -> None:
         """test gVisor runtime verification using a mock docker client"""
@@ -64,9 +64,9 @@ class TestDockerUtility:
 
     def test_socket_read_timeout(self) -> None:
         """Test socket read timeout."""
-        docker = DockerWrapper(image='python', command='python')
+        docker = DockerWrapper(image='python', default_command=['python'])
         # this query should fail as python needs to be started with python3 -i
-        output = docker.exec_run("test query")
+        output = docker.exec_run("test query", timeout=1)
         assert output == "ERROR: timeout"
 
 def test_get_image_template() -> None: