test(learning): add unit tests for learning module

gpt_engineer/applications/cli/learning.py

@@ -88,6 +88,8 @@ class Learning:
  + "(ncertain): "
 )
 
+VALID_INPUTS = ("y", "n", "u")
+
 
 def human_review_input() -> Optional[Review]:
  """
@@ -108,30 +110,30 @@ def human_review_input() -> Optional[Review]:
  print()
 
  ran = input("Did the generated code run at all? " + TERM_CHOICES)
- while ran not in ("y", "n", "u"):
- ran = input("Invalid input. Please enter y, n, or u: ")
-
- perfect = ""
- useful = ""
+ ran = ask_for_valid_input(ran, VALID_INPUTS)
 
  if ran == "y":
  perfect = input(
  "Did the generated code do everything you wanted? " + TERM_CHOICES
  )
- while perfect not in ("y", "n", "u"):
- perfect = input("Invalid input. Please enter y, n, or u: ")
+ perfect = ask_for_valid_input(perfect, VALID_INPUTS)
 
  if perfect != "y":
  useful = input("Did the generated code do anything useful? " + TERM_CHOICES)
- while useful not in ("y", "n", "u"):
- useful = input("Invalid input. Please enter y, n, or u: ")
+ useful = ask_for_valid_input(useful, VALID_INPUTS)
+ else:
+ useful = ""
+ else:
+ perfect = ""
+ useful = ""
 
- comments = ""
  if perfect != "y":
  comments = input(
  "If you have time, please explain what was not working "
  + colored("(ok to leave blank)\n", "light_green")
  )
+ else:
+ comments = ""
 
  return Review(
  raw=", ".join([ran, perfect, useful]),
@@ -142,6 +144,12 @@ def human_review_input() -> Optional[Review]:
  )
 
 
+def ask_for_valid_input(ran, valid_inputs):
+ while ran not in valid_inputs:
+ ran = input("Invalid input. Please enter y, n, or u: ")
+ return ran
+
+
 def check_collection_consent() -> bool:
  """
  Check if the user has given consent to store their data.
@@ -196,23 +204,25 @@ def extract_learning(
 
  Parameters
  ----------
+ prompt : str
+ The prompt used.
  model : str
  The name of the model used.
  temperature : float
  The temperature used.
- steps : List[Step]
+ config : List[Step]
  The list of steps.
- dbs : DBs
- The databases containing the input, logs, memory, and workspace.
- steps_file_hash : str
+ memory : DiskMemory
+ The class containing the input, logs, memory, and workspace.
+ review : Review
  The hash of the steps file.
 
  Returns
  -------
  Learning
  The extracted learning data.
  """
- learning = Learning(
+ return Learning(
  prompt=prompt,
  model=model,
  temperature=temperature,
@@ -221,7 +231,6 @@ def extract_learning(
  logs=memory.to_json(),
  review=review,
  )
- return learning
 
 
 def get_session() -> str:

tests/applications/cli/test_learning.py

@@ -0,0 +1,107 @@
+from unittest import mock
+
+from gpt_engineer.applications.cli import learning
+from gpt_engineer.applications.cli.learning import Learning
+from gpt_engineer.core.default.disk_memory import DiskMemory
+
+
+def test_human_review_input_no_concent_returns_none():
+ with mock.patch.object(learning, "check_collection_consent", return_value=False):
+ result = learning.human_review_input()
+
+ assert result is None
+
+
+def test_human_review_input_consent_code_ran_no_comments():
+ with mock.patch.object(
+ learning, "check_collection_consent", return_value=True
+ ), mock.patch("builtins.input", return_value="y"):
+ result = learning.human_review_input()
+
+ assert result.raw == "y, y, "
+ assert result.ran is True
+ assert result.works is None
+ assert result.comments == ""
+
+
+def test_human_review_input_consent_code_ran_not_perfect_but_useful_no_comments():
+ with mock.patch.object(
+ learning, "check_collection_consent", return_value=True
+ ), mock.patch("builtins.input", side_effect=["y", "n", "y", ""]):
+ result = learning.human_review_input()
+
+ assert result.raw == "y, n, y"
+ assert result.ran is True
+ assert result.works is True
+ assert result.comments == ""
+
+
+def test_check_collection_consent_yes():
+ gpte_consent_mock = mock.Mock()
+ gpte_consent_mock.exists.return_value = True
+ gpte_consent_mock.read_text.return_value = "true"
+
+ with mock.patch.object(learning, "Path", return_value=gpte_consent_mock):
+ result = learning.check_collection_consent()
+
+ assert result is True
+
+
+def test_check_collection_consent_no_ask_collection_consent():
+ with mock.patch.object(learning, "Path") as gpte_consent_mock:
+ gpte_consent_mock.exists.return_value = True
+ gpte_consent_mock.read_text.return_value = "false"
+
+ with mock.patch.object(learning, "ask_collection_consent", return_value=True):
+ result = learning.check_collection_consent()
+
+ assert result is True
+
+
+def test_ask_collection_consent_yes():
+ with mock.patch("builtins.input", return_value="y"):
+ result = learning.ask_collection_consent()
+
+ assert result is True
+
+
+def test_ask_collection_consent_no():
+ with mock.patch("builtins.input", return_value="n"):
+ result = learning.ask_collection_consent()
+
+ assert result is False
+
+
+def test_extract_learning():
+ review = learning.Review(
+ raw="y, n, y",
+ ran=True,
+ works=True,
+ perfect=False,
+ comments="The code is not perfect",
+ )
+ memory = mock.Mock(spec=DiskMemory)
+ memory.to_json.return_value = {"prompt": "prompt"}
+
+ result = learning.extract_learning(
+ "prompt",
+ "model_name",
+ 0.01,
+ ("prompt_tokens", "completion_tokens"),
+ memory,
+ review,
+ )
+
+ assert isinstance(result, Learning)
+
+
+def test_get_session():
+ with mock.patch.object(learning, "Path") as path_mock:
+ # can be better tested with pyfakefs.
+ path_mock.return_value.__truediv__.return_value.exists.return_value = False
+
+ with mock.patch.object(learning, "random") as random_mock:
+ random_mock.randint.return_value = 42
+ result = learning.get_session()
+
+ assert result == "42"