From b68e2e3360fc4c886e47d8f188dc7b432b90ac21 Mon Sep 17 00:00:00 2001 From: Joe Martinez Date: Fri, 5 Jul 2024 00:53:35 -0500 Subject: [PATCH] test_vsearch --- requirements.txt | 108 +++++++++++++++++++++++++++++++++++- tests/test_vector_search.py | 31 ----------- tests/test_vsearch.py | 90 ++++++++++++++++++++++++++++++ 3 files changed, 195 insertions(+), 34 deletions(-) delete mode 100644 tests/test_vector_search.py create mode 100644 tests/test_vsearch.py diff --git a/requirements.txt b/requirements.txt index d2e0ffd..6e3999a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,110 @@ +annotated-types==0.7.0 +anyio==4.4.0 +asgiref==3.8.1 +backoff==2.2.1 +bcrypt==4.1.3 +beautifulsoup4==4.12.3 +Brotli==1.1.0 +build==1.2.1 +cachetools==5.3.3 +certifi==2024.6.2 +charset-normalizer==3.3.2 +chroma-hnswlib==0.7.3 +chromadb==0.5.2 click==8.1.7 +click-default-group==1.2.4 +coloredlogs==15.0.1 +Deprecated==1.2.14 +distro==1.9.0 +dnspython==2.6.1 +email_validator==2.2.0 +fastapi==0.111.0 +fastapi-cli==0.0.4 +filelock==3.15.4 +flatbuffers==24.3.25 +fsspec==2024.6.0 +google-auth==2.30.0 +googleapis-common-protos==1.63.2 +grpcio==1.64.1 +h11==0.14.0 +httpcore==1.0.5 +httptools==0.6.1 +httpx==0.27.0 +huggingface-hub==0.23.4 +humanfriendly==10.0 +idna==3.7 +importlib_metadata==7.1.0 +importlib_resources==6.4.0 +iniconfig==2.0.0 +Jinja2==3.1.4 +kubernetes==30.1.0 +markdown-it-py==3.0.0 +MarkupSafe==2.1.5 +mdurl==0.1.2 +mmh3==4.1.0 +monotonic==1.6 +mpmath==1.3.0 +mutagen==1.47.0 +numpy==1.26.4 +oauthlib==3.2.2 +onnxruntime==1.18.0 openai==1.35.3 -chromadb==0.5.2 +opentelemetry-api==1.25.0 +opentelemetry-exporter-otlp-proto-common==1.25.0 +opentelemetry-exporter-otlp-proto-grpc==1.25.0 +opentelemetry-instrumentation==0.46b0 +opentelemetry-instrumentation-asgi==0.46b0 +opentelemetry-instrumentation-fastapi==0.46b0 +opentelemetry-proto==1.25.0 +opentelemetry-sdk==1.25.0 +opentelemetry-semantic-conventions==0.46b0 +opentelemetry-util-http==0.46b0 +orjson==3.10.5 +overrides==7.7.0 +packaging==24.1 +pluggy==1.5.0 +posthog==3.5.0 +protobuf==4.25.3 +pyasn1==0.6.0 +pyasn1_modules==0.4.0 +pycryptodomex==3.20.0 +pydantic==2.7.4 +pydantic_core==2.18.4 +Pygments==2.18.0 +PyPika==0.48.9 +pyproject_hooks==1.1.0 +pytest==8.2.2 +pytest-order==1.2.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.0.1 +python-multipart==0.0.9 +PyYAML==6.0.1 requests==2.31.0 +requests-oauthlib==2.0.0 rich==13.7.1 +rsa==4.9 +setuptools==70.1.1 +shellingham==1.5.4 +six==1.16.0 +sniffio==1.3.1 +soupsieve==2.5 +sqlite-fts4==1.0.3 sqlite-utils==3.36 -beautifulsoup4==4.12.3 -yt-dlp==2024.5.27 \ No newline at end of file +starlette==0.37.2 +sympy==1.12.1 +tabulate==0.9.0 +tenacity==8.4.2 +tokenizers==0.19.1 +tqdm==4.66.4 +typer==0.12.3 +typing_extensions==4.12.2 +ujson==5.10.0 +urllib3==2.2.2 +uvicorn==0.30.1 +uvloop==0.19.0 +watchfiles==0.22.0 +websocket-client==1.8.0 +websockets==12.0 +wrapt==1.16.0 +yt-dlp==2024.5.27 +zipp==3.19.2 diff --git a/tests/test_vector_search.py b/tests/test_vector_search.py deleted file mode 100644 index cf92e35..0000000 --- a/tests/test_vector_search.py +++ /dev/null @@ -1,31 +0,0 @@ -import unittest -from openai import OpenAI -import os -from yt_fts.vector_search import search_chroma_db - -class TestSearchChromaDb(unittest.TestCase): - def test_search_chroma_db(self): - # Setup - text = "nural network" - scope = "all" - video_id = None - channel_id = None - limit = 5 - client = OpenAI(api_key=os.environ["OPENAI_API_KEY"]) - - result = search_chroma_db(text, scope, channel_id, video_id, limit, client) - self.assertEqual(len(result), 5) - - def test_search_chroma_db_channel(self): - text = "nural network" - scope = "channel" - video_id = None - channel_id = "1" - limit = 5 - client = OpenAI(api_key=os.environ["OPENAI_API_KEY"]) - - result = search_chroma_db(text, scope, channel_id, video_id, limit, client) - self.assertEqual(len(result), 5) - -if __name__ == '__main__': - unittest.main() \ No newline at end of file diff --git a/tests/test_vsearch.py b/tests/test_vsearch.py new file mode 100644 index 0000000..a1b8913 --- /dev/null +++ b/tests/test_vsearch.py @@ -0,0 +1,90 @@ +import pytest +import requests +import sqlite3 +import os +import shutil +import zipfile +import sys +from io import BytesIO +from click.testing import CliRunner +from yt_fts.yt_fts import download, cli + +CONFIG_DIR = os.path.expanduser('~/.config/yt-fts') + +@pytest.fixture +def runner(): + return CliRunner() + + +def reset_testing_env(): + if os.path.exists(CONFIG_DIR): + if os.environ.get('YT_FTS_TEST_RESET', 'true').lower() == 'true': + shutil.rmtree(CONFIG_DIR) + fetch_and_unzip_test_db() + else: + print('running tests with existing db') + + +def fetch_and_unzip_test_db(): + url = "https://yt-fts-testdb-server.notjoemartinez.workers.dev/yt-fts/test_dbs/2024-07-04.zip" + response = requests.get(url) + if response.status_code == 200: + with zipfile.ZipFile(BytesIO(response.content)) as zip_ref: + for file in zip_ref.namelist(): + if file.startswith('2024-07-04/'): + extracted_path = zip_ref.extract(file, CONFIG_DIR) + print("extracted path: ", extracted_path) + + new_path = os.path.join(CONFIG_DIR, os.path.relpath(extracted_path, os.path.join(CONFIG_DIR, '2024-07-04'))) + print(f"new_path: {new_path}") + os.makedirs(os.path.dirname(new_path), exist_ok=True) + + if os.path.exists(new_path): + if os.path.isdir(new_path): + for item in os.listdir(extracted_path): + s = os.path.join(extracted_path, item) + d = os.path.join(new_path, item) + if os.path.isdir(s): + shutil.copytree(s, d, dirs_exist_ok=True) + else: + shutil.copy2(s, d) + else: + os.remove(new_path) + shutil.move(extracted_path, new_path) + else: + shutil.move(extracted_path, new_path) + + shutil.rmtree(os.path.join(CONFIG_DIR, '2024-07-04'), ignore_errors=True) + + else: + raise Exception(f"Failed to download the zip file. Status code: {response.status_code}") + + print(f"Successfully extracted") + + +def get_test_db(): + conn = sqlite3.connect(f"{CONFIG_DIR}/subtitles.db") + curr = conn.cursor() + return curr + + +def test_vsearch(runner, capsys): + reset_testing_env() + + result = runner.invoke(cli, [ + 'vsearch', + '-c', + '3', + 'icbm gambit', + ]) + + print(result.output) + captured = capsys.readouterr() + output = captured.out + + assert "Title: Intercontinental Ballistic Missile Gambit (real opening)" in output + assert "missile attack that will leave your opponent's position in unorganized chaos" in output + + +if __name__ == "__main__": + pytest.main([__file__]) \ No newline at end of file