Merge pull request wise-agents#375 from kabir/tests_use_existing_env_…

…vars

wise-agents#251 Don't set env vars in tests.

tests/wiseagents/__init__.py

@@ -1,8 +1,11 @@
 # This is the __init__.py file for the wiseagents package
 
 # Import any modules or subpackages here
+from .testing_utils import assert_env_var_set, assert_env_vars_set, assert_standard_variables_set
 
 # Define any necessary initialization code here
 
 # Optionally, you can define __all__ to specify the public interface of the package
-# __all__ = ['module1', 'module2', 'subpackage']
+# __all__ = ['module1', 'module2', 'subpackage']
+
+__all__ = ['assert_env_var_set', 'assert_env_vars_set', 'assert_standard_variables_set']

tests/wiseagents/agents/test_graph_rag_challenger.py

@@ -11,6 +11,7 @@
 from wiseagents.graphdb import Entity, GraphDocument, Neo4jLangChainWiseAgentGraphDB, Relationship, Source
 from wiseagents.llm import OpenaiAPIWiseAgentLLM
 from wiseagents.transports import StompWiseAgentTransport
+from tests.wiseagents import assert_standard_variables_set
 
 cond = threading.Condition()
 assertError : AssertionError = None
@@ -19,11 +20,10 @@
 collection_name = "test-vector-db"
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
+
+    assert_standard_variables_set()
+
     # Ensure that nothing exists the graph DB
-    original_neo4j_username = os.environ.get("NEO4J_USERNAME")
-    os.environ["NEO4J_USERNAME"] = "neo4j"
-    original_neo4j_password = os.environ.get("NEO4J_PASSWORD")
-    os.environ["NEO4J_PASSWORD"] = "neo4jpassword"
     graph_db = Neo4jLangChainWiseAgentGraphDB(properties=["name", "type"], collection_name=collection_name,
                                               url="bolt://localhost:7687", refresh_graph_schema=False)
     graph_db.query("MATCH (n)-[r]-() DELETE r")
@@ -54,33 +54,6 @@ def run_after_all_tests():
     graph_db.delete_vector_db()
     graph_db.close()
 
-    # Clean up environment variables
-    if original_neo4j_username is None:
-        del os.environ["NEO4J_USERNAME"]
-    else:
-        os.environ["NEO4J_USERNAME"] = original_neo4j_username
-    if original_neo4j_password is None:
-        del os.environ["NEO4J_PASSWORD"]
-    else:
-        os.environ["NEO4J_PASSWORD"] = original_neo4j_password
-
-
-
-
-
-
-def set_env_variable(env_variable: str, value: str) -> str:
-    original_value = os.environ.get(env_variable)
-    os.environ[env_variable] = value
-    return original_value
-
-
-def reset_env_variable(env_variable: str, original_value: str):
-    if original_value is None:
-        del os.environ[env_variable]
-    else:
-        os.environ[env_variable] = original_value
-
 
 def get_connection_string():
     return f"postgresql+psycopg://{os.environ['POSTGRES_USER']}:{os.environ['POSTGRES_PASSWORD']}@localhost:6024/{os.environ['POSTGRES_DB']}"

tests/wiseagents/agents/test_phased_coordinator.py

@@ -9,12 +9,14 @@
     PhasedCoordinatorWiseAgent
 from wiseagents.llm import OpenaiAPIWiseAgentLLM
 from wiseagents.transports import StompWiseAgentTransport
+from tests.wiseagents import assert_standard_variables_set
 
 cond = threading.Condition()
 assertError : AssertionError = None
 
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
+    assert_standard_variables_set()
     yield
 
 
@@ -37,6 +39,7 @@ def test_phased_coordinator():
     """
     try:
         groq_api_key = os.getenv("GROQ_API_KEY")
+
         global assertError
 
         llm = OpenaiAPIWiseAgentLLM(system_message="You are a helpful assistant.",

tests/wiseagents/agents/test_rag_challenger.py

@@ -9,19 +9,15 @@
 from wiseagents.llm import OpenaiAPIWiseAgentLLM
 from wiseagents.transports import StompWiseAgentTransport
 from wiseagents.vectordb import Document, PGVectorLangChainWiseAgentVectorDB
+from tests.wiseagents import assert_standard_variables_set
 
 cond = threading.Condition()
 assertError : AssertionError = None
 
 
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
-    original_postgres_user = set_env_variable("POSTGRES_USER", "postgres")
-    original_postgres_password = set_env_variable("POSTGRES_PASSWORD", "postgres")
-    original_postgres_db = set_env_variable("POSTGRES_DB", "postgres")
-    original_stomp_user = set_env_variable("STOMP_USER", "artemis")
-    original_stomp_password = set_env_variable("STOMP_PASSWORD", "artemis")
-
+    assert_standard_variables_set()
     # Vector DB set up
     pg_vector_db = PGVectorLangChainWiseAgentVectorDB(get_connection_string())
     pg_vector_db.delete_collection("wise-agents-collection")
@@ -40,29 +36,6 @@ def run_after_all_tests():
     # Vector DB clean up
     pg_vector_db.delete_collection("wise-agents-collection")
 
-    # Clean up environment variables
-    reset_env_variable("POSTGRES_USER", original_postgres_user)
-    reset_env_variable("POSTGRES_PASSWORD", original_postgres_password)
-    reset_env_variable("POSTGRES_DB", original_postgres_db)
-    reset_env_variable("STOMP_USER", original_stomp_user)
-    reset_env_variable("STOMP_PASSWORD", original_stomp_password)
-
-
-
-
-
-def set_env_variable(env_variable: str, value: str) -> str:
-    original_value = os.environ.get(env_variable)
-    os.environ[env_variable] = value
-    return original_value
-
-
-def reset_env_variable(env_variable: str, original_value: str):
-    if original_value is None:
-        del os.environ[env_variable]
-    else:
-        os.environ[env_variable] = original_value
-
 
 def get_connection_string():
     return f"postgresql+psycopg://{os.environ['POSTGRES_USER']}:{os.environ['POSTGRES_PASSWORD']}@localhost:6024/{os.environ['POSTGRES_DB']}"

tests/wiseagents/agents/test_sequential_coordinator.py

@@ -12,6 +12,7 @@
     SequentialMemoryCoordinatorWiseAgent
 from wiseagents.llm import OpenaiAPIWiseAgentLLM, WiseAgentLLM
 from wiseagents.transports import StompWiseAgentTransport
+from tests.wiseagents import assert_standard_variables_set
 
 cond1 = threading.Condition()
 cond2 = threading.Condition()
@@ -20,6 +21,7 @@
 
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
+    assert_standard_variables_set()
     yield
 
 

tests/wiseagents/agents/test_tools.py

@@ -1,6 +1,5 @@
 import json
 import logging
-import os
 import threading
 
 import pytest
@@ -9,10 +8,12 @@
 from wiseagents.agents import LLMWiseAgentWithTools, PassThroughClientAgent
 from wiseagents.llm import OpenaiAPIWiseAgentLLM
 from wiseagents.transports.stomp import StompWiseAgentTransport
+from tests.wiseagents import assert_standard_variables_set
 
 
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
+    assert_standard_variables_set()
     yield
 
 

tests/wiseagents/agents/test_tools_fakeLLM.py

@@ -10,10 +10,12 @@
 from wiseagents.agents import LLMWiseAgentWithTools, PassThroughClientAgent
 from wiseagents.llm import WiseAgentRemoteLLM
 from wiseagents.transports.stomp import StompWiseAgentTransport
+from tests.wiseagents import assert_standard_variables_set
 
 
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
+    assert_standard_variables_set()
     yield
 
 

tests/wiseagents/agents/test_tools_groq.py

@@ -9,10 +9,12 @@
 from wiseagents.agents import LLMWiseAgentWithTools, PassThroughClientAgent
 from wiseagents.llm import OpenaiAPIWiseAgentLLM
 from wiseagents.transports.stomp import StompWiseAgentTransport
+from tests.wiseagents import assert_standard_variables_set
 
 
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
+    assert_standard_variables_set()
     yield
 
 

tests/wiseagents/cli/test_WiseAgentCli.py

@@ -7,17 +7,11 @@
 from wiseagents.cli.wise_agent_cli import main
 from wiseagents.graphdb import Entity, GraphDocument, Neo4jLangChainWiseAgentGraphDB, Relationship, Source
 from wiseagents.vectordb import Document, PGVectorLangChainWiseAgentVectorDB
-
+from tests.wiseagents import assert_standard_variables_set
 
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
-    original_postgres_user = set_env_variable("POSTGRES_USER", "postgres")
-    original_postgres_password = set_env_variable("POSTGRES_PASSWORD", "postgres")
-    original_postgres_db = set_env_variable("POSTGRES_DB", "postgres")
-    original_stomp_user = set_env_variable("STOMP_USER", "artemis")
-    original_stomp_password = set_env_variable("STOMP_PASSWORD", "artemis")
-    original_neo4j_username = set_env_variable("NEO4J_USERNAME", "neo4j")
-    original_neo4j_password = set_env_variable("NEO4J_PASSWORD", "neo4jpassword")
+    assert_standard_variables_set()
 
     # Vector DB set up
     pg_vector_db = PGVectorLangChainWiseAgentVectorDB(get_connection_string())
@@ -60,45 +54,11 @@ def run_after_all_tests():
     graph_db.delete_vector_db()
     graph_db.close()
 
-    # Clean up environment variables
-    reset_env_variable("POSTGRES_USER", original_postgres_user)
-    reset_env_variable("POSTGRES_PASSWORD", original_postgres_password)
-    reset_env_variable("POSTGRES_DB", original_postgres_db)
-    reset_env_variable("STOMP_USER", original_stomp_user)
-    reset_env_variable("STOMP_PASSWORD", original_stomp_password)
-    reset_env_variable("NEO4J_USERNAME", original_neo4j_username)
-    reset_env_variable("NEO4J_PASSWORD", original_neo4j_password)
-
-
-
-
-def set_env_variable(env_variable: str, value: str) -> str:
-    original_value = os.environ.get(env_variable)
-    os.environ[env_variable] = value
-    return original_value
-
-
-def reset_env_variable(env_variable: str, original_value: str):
-    if original_value is None:
-        del os.environ[env_variable]
-    else:
-        os.environ[env_variable] = original_value
-
 
 def get_connection_string():
     return f"postgresql+psycopg://{os.environ['POSTGRES_USER']}:{os.environ['POSTGRES_PASSWORD']}@localhost:6024/{os.environ['POSTGRES_DB']}"
 
 
-def set_env(monkeypatch):
-    """
-        This test requires a running pgvector instance. The required
-        vector database can be started using the run_vectordb.sh script.
-    """
-    monkeypatch.setenv("POSTGRES_USER", "postgres")
-    monkeypatch.setenv("POSTGRES_PASSWORD", "postgres")
-    monkeypatch.setenv("POSTGRES_DB", "postgres")
-
-
 @pytest.mark.needsllm
 def test_cli_with_rag_agent(monkeypatch, pytestconfig, capsys):
     inputs = ['/load-agents', '', '/send', 'RAGWiseAgent1', 'Who won the NBA championship in 2024?', '/exit']

tests/wiseagents/graphdb/test_Neo4jLangChainWiseAgentGraphDB.py

@@ -1,17 +1,13 @@
-import os
-
 import pytest
 
 from wiseagents.graphdb import Entity, GraphDocument, Neo4jLangChainWiseAgentGraphDB, Relationship, Source
+from tests.wiseagents import assert_standard_variables_set
 
 collection_name = "test-vector-db"
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
+    assert_standard_variables_set()
     # Ensure that nothing exists the graph DB
-    original_neo4j_username = os.environ.get("NEO4J_USERNAME")
-    os.environ["NEO4J_USERNAME"] = "neo4j"
-    original_neo4j_password = os.environ.get("NEO4J_PASSWORD")
-    os.environ["NEO4J_PASSWORD"] = "neo4jpassword"
     graph_db = Neo4jLangChainWiseAgentGraphDB(properties=["name", "type"], collection_name=collection_name,
                                               url="bolt://localhost:7687", refresh_graph_schema=False)
     graph_db.query("MATCH (n)-[r]-() DELETE r")
@@ -27,30 +23,9 @@ def run_after_all_tests():
     graph_db.delete_vector_db()
     graph_db.close()
 
-    # Clean up environment variables
-    if original_neo4j_username is None:
-        del os.environ["NEO4J_USERNAME"]
-    else:
-        os.environ["NEO4J_USERNAME"] = original_neo4j_username
-    if original_neo4j_password is None:
-        del os.environ["NEO4J_PASSWORD"]
-    else:
-        os.environ["NEO4J_PASSWORD"] = original_neo4j_password
-
-
-def set_env(monkeypatch):
-    """
-        This test requires a running Neo4j instance on bolt://localhost:7687. The required
-        graph database can be started using the run_graphdb.sh script.
-    """
-    monkeypatch.setenv("NEO4J_USERNAME", "neo4j")
-    assert os.environ.get("NEO4J_USERNAME") == "neo4j"
-    monkeypatch.setenv("NEO4J_PASSWORD", "neo4jpassword")
-    assert os.environ.get("NEO4J_PASSWORD") == "neo4jpassword"
-
-
-def test_insert_graph_documents_and_query(monkeypatch):
-    set_env(monkeypatch)
+
+
+def test_insert_graph_documents_and_query():
     graph_db = Neo4jLangChainWiseAgentGraphDB(properties=["name", "type"], collection_name=collection_name,
                                               url="bolt://localhost:7687", refresh_graph_schema=False)
 
@@ -96,8 +71,7 @@ def test_insert_graph_documents_and_query(monkeypatch):
 
 
 
-def test_insert_entity_and_query(monkeypatch):
-    set_env(monkeypatch)
+def test_insert_entity_and_query():
     graph_db = Neo4jLangChainWiseAgentGraphDB(properties=["name", "type"], collection_name=collection_name,
                                               url="bolt://localhost:7687", refresh_graph_schema=False)
 
@@ -115,8 +89,7 @@ def test_insert_entity_and_query(monkeypatch):
         graph_db.close()
 
 
-def test_insert_relationship_and_query(monkeypatch):
-    set_env(monkeypatch)
+def test_insert_relationship_and_query():
     graph_db = Neo4jLangChainWiseAgentGraphDB(properties=["name", "type"], collection_name=collection_name,
                                               url="bolt://localhost:7687", refresh_graph_schema=False)
 

tests/wiseagents/llm/test_WiseAgentRemoteLLM.py

@@ -1,11 +1,11 @@
 import pytest
 
-from wiseagents import WiseAgentRegistry
 from wiseagents.llm import OpenaiAPIWiseAgentLLM
-
+from tests.wiseagents import assert_standard_variables_set
 
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
+    assert_standard_variables_set()
     yield
 
 

tests/wiseagents/test_WiseAgentRegistry.py

@@ -4,10 +4,11 @@
 
 from wiseagents import WiseAgent, WiseAgentContext, WiseAgentMessage, WiseAgentMetaData, WiseAgentRegistry, WiseAgentTransport
 from wiseagents.transports.stomp import StompWiseAgentTransport
-
+from tests.wiseagents import assert_standard_variables_set
 
 @pytest.fixture(scope="session", autouse=True)
 def run_after_all_tests():
+    assert_standard_variables_set()
     yield