"""Test cases for procedures defined in llms """ # pylint: disable=C0301 # pylint: disable=W0613 from datetime import datetime, timedelta from unittest.mock import patch from unittest.mock import MagicMock import pytest from gn3.llms.process import fetch_pubmed from gn3.llms.process import parse_context from gn3.llms.process import format_bibliography_info from gn3.llms.errors import LLMError from gn3.api.llm import clean_query from gn3.api.llm import is_verified_anonymous_user from gn3.api.llm import is_valid_address from gn3.api.llm import check_rate_limiter FAKE_NOW = datetime(2025, 1, 1, 12, 0, 0) @pytest.mark.unit_test def test_parse_context(): """test for parsing doc id context""" def mock_get_info(doc_id): return f"Info for {doc_id}" def mock_format_bib(doc_info): return f"Formatted Bibliography: {doc_info}" parsed_result = parse_context({ "doc1": [{"text": "Summary 1"}, {"text": "Summary 2"}], "doc2": [{"text": "Summary 3"}, {"text": "Summary 4"}], }, mock_get_info, mock_format_bib) expected_result = [ { "doc_id": "doc1", "bibInfo": "Formatted Bibliography: Info for doc1", "comboTxt": "\tSummary 1\tSummary 2", }, { "doc_id": "doc2", "bibInfo": "Formatted Bibliography: Info for doc2", "comboTxt": "\tSummary 3\tSummary 4", }, ] assert parsed_result == expected_result @pytest.mark.unit_test def test_format_bib_info(): """Test for formatting bibliography info """ mock_fahamu_bib_info = [ { "author": "J.m", "firstName": "john", "title": "Genes and aging", "year": 2013, "doi": "https://Articles.com/12231" }, "2019-Roy-Evaluation of Sirtuin-3 probe quality and co-expressed genes", "2015 - Differential regional and cellular distribution of TFF3 peptide in the human brain.txt"] expected_result = [ "J.m.Genes and aging.2013.https://Articles.com/12231 ", "2019-Roy-Evaluation of Sirtuin-3 probe quality and co-expressed genes", "2015 - Differential regional and cellular distribution of TFF3 peptide in the human brain" ] assert all((format_bibliography_info(data) == expected for data, expected in zip(mock_fahamu_bib_info, expected_result))) @pytest.mark.unit_test def test_fetching_pubmed_info(monkeypatch): """Test for fetching and populating pubmed data with pubmed info""" def mock_load_file(_filename, _dir_path): return { "12121": { "Abstract": "items1", "Author": "A1" } } # patch the module with the mocked function monkeypatch.setattr("gn3.llms.process.load_file", mock_load_file) expected_results = [ { "title": "Genes", "year": "2014", "doi": "https/article/genes/12121", "doc_id": "12121", "pubmed": { "Abstract": "items1", "Author": "A1" } }, { "title": "Aging", "year": "2014", "doc_id": "12122" } ] data = [{ "title": "Genes", "year": "2014", "doi": "https/article/genes/12121", "doc_id": "12121", }, { "title": "Aging", "year": "2014", "doc_id": "12122" }] assert (fetch_pubmed(data, "/pubmed.json", "data/") == expected_results) @pytest.mark.unit_test def test_clean_query(): """Test function for cleaning up query""" assert clean_query("!what is genetics.") == "what is genetics" assert clean_query("hello test?") == "hello test" assert clean_query(" hello test with space?") == "hello test with space" @pytest.mark.unit_test def test_is_verified_anonymous_user(): """Test function for verifying anonymous user metadata""" assert is_verified_anonymous_user({}) is False assert is_verified_anonymous_user({"Anonymous-Id" : "qws2121dwsdwdwe", "Anonymous-Status" : "verified"}) is True @pytest.mark.unit_test def test_is_valid_address() : """Test function checks if is a valid ip address is valid""" assert is_valid_address("invalid_ip") is False assert is_valid_address("127.0.0.1") is True @patch("gn3.api.llm.datetime") @patch("gn3.api.llm.db.connection") @patch("gn3.api.llm.is_valid_address", return_value=True) @pytest.mark.unit_test def test_first_time_visitor(mock_is_valid, mock_db_conn, mock_datetime): """Test rate limiting for first-time visitor""" mock_datetime.utcnow.return_value = FAKE_NOW mock_datetime.strptime = datetime.strptime # keep real one mock_datetime.strftime = datetime.strftime # keep real one # Set up DB mock mock_conn = MagicMock() mock_cursor = MagicMock() mock_conn.__enter__.return_value = mock_conn mock_conn.cursor.return_value = mock_cursor mock_cursor.fetchone.return_value = None mock_db_conn.return_value = mock_conn result = check_rate_limiter("127.0.0.1", "test/llm.db", "Chromosome x") assert result is True mock_cursor.execute.assert_any_call(""" INSERT INTO Limiter(identifier, tokens, expiry_time) VALUES (?, ?, ?) """, ("127.0.0.1", 4, "2025-01-01 12:24:00")) @patch("gn3.api.llm.datetime") @patch("gn3.api.llm.db.connection") @patch("gn3.api.llm.is_valid_address", return_value=True) @pytest.mark.unit_test def test_visitor_at_limit(mock_is_valid, mock_db_conn, mock_datetime): """Test rate limiting for Visitor at limit""" mock_datetime.utcnow.return_value = FAKE_NOW mock_datetime.strptime = datetime.strptime # keep real one mock_datetime.strftime = datetime.strftime mock_conn = MagicMock() mock_cursor = MagicMock() mock_conn.__enter__.return_value = mock_conn mock_conn.cursor.return_value = mock_cursor fake_expiry = (FAKE_NOW + timedelta(minutes=10)).strftime("%Y-%m-%d %H:%M:%S") mock_cursor.fetchone.return_value = (0, fake_expiry) #token returned are 0 mock_db_conn.return_value = mock_conn with pytest.raises(LLMError) as exc_info: check_rate_limiter("127.0.0.1", "test/llm.db", "Chromosome x") # assert llm error with correct message is raised assert exc_info.value.args == ('Rate limit exceeded. Please try again later.', 'Chromosome x') @patch("gn3.api.llm.datetime") @patch("gn3.api.llm.db.connection") @patch("gn3.api.llm.is_valid_address", return_value=True) @pytest.mark.unit_test def test_visitor_with_tokens(mock_is_valid, mock_db_conn, mock_datetime): """Test rate limiting for user with valid tokens""" mock_datetime.utcnow.return_value = FAKE_NOW mock_datetime.strptime = datetime.strptime # Use real versions mock_datetime.strftime = datetime.strftime mock_conn = MagicMock() mock_cursor = MagicMock() mock_conn.__enter__.return_value = mock_conn mock_conn.cursor.return_value = mock_cursor fake_expiry = (FAKE_NOW + timedelta(minutes=10)).strftime("%Y-%m-%d %H:%M:%S") mock_cursor.fetchone.return_value = (3, fake_expiry) # Simulate 3 tokens mock_db_conn.return_value = mock_conn results = check_rate_limiter("127.0.0.1", "test/llm.db", "Chromosome x") assert results is True mock_cursor.execute.assert_any_call(""" UPDATE Limiter SET tokens = tokens - 1 WHERE identifier = ? AND tokens > 0 """, ("127.0.0.1",)) @patch("gn3.api.llm.datetime") @patch("gn3.api.llm.db.connection") @patch("gn3.api.llm.is_valid_address", return_value=True) @pytest.mark.unit_test def test_visitor_token_expired(mock_is_valid, mock_db_conn, mock_datetime): """Test rate limiting for expired tokens""" mock_datetime.utcnow.return_value = FAKE_NOW mock_datetime.strptime = datetime.strptime mock_datetime.strftime = datetime.strftime mock_conn = MagicMock() mock_cursor = MagicMock() mock_conn.__enter__.return_value = mock_conn mock_conn.cursor.return_value = mock_cursor fake_expiry = (FAKE_NOW - timedelta(minutes=10)).strftime("%Y-%m-%d %H:%M:%S") mock_cursor.fetchone.return_value = (3, fake_expiry) # Simulate 3 tokens mock_db_conn.return_value = mock_conn result = check_rate_limiter("127.0.0.1", "test/llm.db", "Chromosome x") assert result is True mock_cursor.execute.assert_any_call(""" UPDATE Limiter SET tokens = ?, expiry_time = ? WHERE identifier = ? """, (4, "2025-01-01 12:24:00", "127.0.0.1"))