add api key, formatting fixes

integration_tests.py

@@ -5,36 +5,44 @@
 import time
 from main import chain
 
+os.getenv("OPENAI_API_KEY")
+
 known_colors = {
-    'red', 'blue', 'yellow',  
-    'green', 'orange', 'purple', 
+    "red",
+    "blue",
+    "yellow",
+    "green",
+    "orange",
+    "purple",
 }
 
+
 @pytest.mark.integration
 def test_chain_invoke():
     # Set maximum time (in seconds)
-    max_allowed_time = 5  
-    
+    max_allowed_time = 5
+
     # Record the start time
     start_time = time.time()
-    
+
     # Invoke the chain
     response = chain.invoke({"text": "colors"})
-    
+
     # Record the end time
     end_time = time.time()
-    
+
     # Calculate the duration
     duration = end_time - start_time
-    
+
     # Assert that the response time is less than the maximum allowed time
-    assert duration < max_allowed_time, f"API response took {duration} seconds, which is longer than the allowed {max_allowed_time} seconds"
-
+    assert (
+        duration < max_allowed_time
+    ), f"API response took {duration} seconds, which is longer than the allowed {max_allowed_time} seconds"
+
     # Check if the API returns a list of colors:
     assert isinstance(response, list)
     assert len(response) == 5  # assuming API returns 5 colors
 
     # Assert that each color in the response is a known color
     for color in response:
         assert color in known_colors, f"Unknown color: {color}"
-

unit_tests.py

@@ -6,6 +6,8 @@
 import string
 from main import CommaSeparatedListOutputParser, template
 
+os.getenv("OPENAI_API_KEY")
+
 
 @pytest.fixture
 def parser():
@@ -16,22 +18,27 @@ def parser():
 def spell():
     return SpellChecker()
 
+
 @pytest.fixture
 def template_spellcheck(spell):
     # Remove punctuation from the template string and then split it into words
-    translator = str.maketrans('', '', string.punctuation)
+    translator = str.maketrans("", "", string.punctuation)
     no_punctuation_template = template.translate(translator)
     words = no_punctuation_template.split()
-    
+
     misspelled = spell.unknown(words)
-    assert not misspelled, f"The template contains the following misspelled words: {', '.join(misspelled)}"
+    assert (
+        not misspelled
+    ), f"The template contains the following misspelled words: {', '.join(misspelled)}"
+
 
 def test_parse_with_valid_input(parser, spell):
     input_text = "red, blue, green, yellow, purple"
     expected_output = ["red", "blue", "green", "yellow", "purple"]
     assert parser.parse(input_text) == expected_output
     assert not spell.unknown(expected_output)
 
+
 def test_parse_with_empty_input(parser):
     input_text = ""
     expected_output = [""]
@@ -42,13 +49,16 @@ def test_parse_with_single_element(parser, spell):
     input_text = "red"
     expected_output = ["red"]
     assert parser.parse(input_text) == expected_output
-    assert not spell.unknown(expected_output) 
+    assert not spell.unknown(expected_output)
 
 
 def test_spellchecking_colors(spell):
     colors = ["red", "blue", "green", "yellow", "purple"]
     misspelled = spell.unknown(colors)
-    assert not misspelled, f"The following colors are misspelled: {', '.join(misspelled)}"
+    assert (
+        not misspelled
+    ), f"The following colors are misspelled: {', '.join(misspelled)}"
+
 
 def test_template_spelling(template_spellcheck):
     pass