added tests

tests.py

@@ -4,7 +4,8 @@ import shutil
 import pytest as pytest
 from grouped_sampling import GroupedSamplingPipeLine
 
-from hanlde_form_submit import create_pipeline
+from on_server_start import download_useful_models
+from hanlde_form_submit import create_pipeline, on_form_submit
 from prompt_engeneering import rewrite_prompt
 from supported_models import get_supported_model_names
 
@@ -24,10 +25,40 @@ def test_get_supported_model_names():
     supported_model_names = get_supported_model_names()
     assert len(supported_model_names) > 0
     assert "gpt2" in supported_model_names
-    assert all([isinstance(name, str) for name in supported_model_names])
+    assert all(isinstance(name, str) for name in supported_model_names)
 
 
-@pytest.mark.parametrize("model_name", get_supported_model_names())
+def test_on_server_start():
+    if os.path.exists(HUGGING_FACE_CACHE_DIR):
+        shutil.rmtree(HUGGING_FACE_CACHE_DIR)
+    assert not os.path.exists(HUGGING_FACE_CACHE_DIR)
+    download_useful_models()
+    assert os.path.exists(HUGGING_FACE_CACHE_DIR)
+    assert len(os.listdir(HUGGING_FACE_CACHE_DIR)) > 0
+
+
+def test_on_form_submit():
+    model_name = "gpt2"
+    output_length = 10
+    prompt = "Answer yes or no, is the sky blue?"
+    output = on_form_submit(model_name, output_length, prompt)
+    assert output is not None
+    assert len(output) > 0
+    empty_prompt = ""
+    with pytest.raises(ValueError):
+        on_form_submit(model_name, output_length, empty_prompt)
+    unsupported_model_name = "unsupported_model_name"
+    with pytest.raises(ValueError):
+        on_form_submit(unsupported_model_name, output_length, prompt)
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    get_supported_model_names(
+        min_number_of_downloads=1000,
+        min_number_of_likes=100,
+    )
+)
 def test_create_pipeline(model_name: str):
     pipeline: GroupedSamplingPipeLine = create_pipeline(model_name, 5)
     assert pipeline is not None