test: Added tests for library (#45)

* test: Added tests for library

* fix: CR

Author: asafgardin

tests/integration_tests/clients/resources/library_file.txt

@@ -0,0 +1,19 @@
+Albert Einstein was a renowned physicist who made significant contributions to the field of theoretical physics. Born on March 14, 1879, in Ulm, in the Kingdom of Württemberg in the German Empire, Einstein's early life showed signs of his later intellectual prowess.
+
+Einstein attended the Swiss Federal Institute of Technology in Zurich, where he studied physics and mathematics. Despite facing challenges and financial difficulties, he persevered in his studies and graduated in 1900. After graduation, he struggled to secure a teaching position but eventually found work as a patent examiner at the Swiss Patent Office.
+
+In 1905, often referred to as his "miracle year," Einstein published four groundbreaking papers that transformed the scientific landscape. These papers covered the photoelectric effect, Brownian motion, special relativity, and the famous equation E=mc², demonstrating the equivalence of mass and energy.
+
+His theory of special relativity, published in the paper "On the Electrodynamics of Moving Bodies," challenged traditional notions of space and time. It introduced the concept of spacetime and showed that time is relative, depending on the observer's motion.
+
+In 1915, Einstein presented the general theory of relativity, providing a new understanding of gravitation. According to general relativity, massive objects like planets and stars cause a curvature in spacetime, influencing the motion of other objects. This theory successfully explained phenomena like the bending of light around massive objects.
+
+Einstein's work laid the foundation for modern cosmology and astrophysics. His predictions, such as the bending of light by gravity, were later confirmed through experiments and observations.
+
+Apart from his scientific endeavors, Einstein was an advocate for civil rights, pacifism, and Zionism. He spoke out against discrimination and injustice, using his platform to promote social and political causes. In 1933, Einstein fled Nazi Germany and settled in the United States, where he continued his scientific research.
+
+Einstein received the Nobel Prize in Physics in 1921 for his explanation of the photoelectric effect. Despite his immense contributions to science, he remained humble and often expressed a deep curiosity about the mysteries of the universe.
+
+In the latter part of his life, Einstein worked towards a unified field theory, attempting to combine electromagnetism and gravity into a single framework. However, this goal remained elusive, and Einstein's efforts in this direction were not as successful as his earlier work.
+
+Albert Einstein passed away on April 18, 1955, leaving behind a legacy that continues to shape our understanding of the physical world. His intellectual brilliance, coupled with his commitment to social justice, has made him an enduring symbol of scientific achievement and moral responsibility. The impact of Einstein's ideas extends far beyond the realm of physics, influencing fields as diverse as philosophy, literature, and popular culture.

tests/integration_tests/clients/studio/conftest.py

@@ -0,0 +1,45 @@
+import time
+from pathlib import Path
+
+import pytest
+
+from ai21 import AI21Client
+
+LIBRARY_FILE_TO_UPLOAD = str(Path(__file__).parent.parent / "resources" / "library_file.txt")
+DEFAULT_LABELS = ["einstein", "science"]
+
+
+def _wait_for_file_to_process(client: AI21Client, file_id: str, timeout: float = 20):
+    start_time = time.time()
+
+    elapsed_time = time.time() - start_time
+    while elapsed_time < timeout:
+        file_response = client.library.files.get(file_id)
+
+        if file_response.status == "PROCESSED":
+            return
+
+        elapsed_time = time.time() - start_time
+        time.sleep(0.5)
+
+    raise TimeoutError(f"Timeout: {timeout} seconds passed. File processing not completed")
+
+
+def _delete_file(client: AI21Client, file_id: str):
+    _wait_for_file_to_process(client, file_id)
+    client.library.files.delete(file_id)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def file_in_library():
+    """
+    Uploads a file to the library and deletes it after the test is done
+    This happens in a scope of a module so the file is uploaded only once
+    :return: file_id: str
+    """
+    client = AI21Client()
+
+    file_id = client.library.files.create(file_path=LIBRARY_FILE_TO_UPLOAD, labels=DEFAULT_LABELS)
+    _wait_for_file_to_process(client, file_id)
+    yield file_id
+    _delete_file(client, file_id=file_id)

tests/integration_tests/clients/studio/test_library.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+from time import sleep
+
+from ai21 import AI21Client
+from tests.integration_tests.clients.studio.conftest import LIBRARY_FILE_TO_UPLOAD, DEFAULT_LABELS
+
+
+def test_library__when_upload__should_get_file_id(file_in_library: str):
+    assert file_in_library is not None
+
+
+def test_library__when_list__should_get_file_id_in_list_of_files(file_in_library: str):
+    client = AI21Client()
+
+    files = client.library.files.list()
+    assert files[0].file_id == file_in_library
+    assert files[0].name == Path(LIBRARY_FILE_TO_UPLOAD).name
+
+
+def test_library__when_get__should_match_file_id(file_in_library: str):
+    client = AI21Client()
+
+    file_response = client.library.files.get(file_in_library)
+    assert file_response.file_id == file_in_library
+
+
+def test_library__when_update__should_update_labels_successfully(file_in_library: str):
+    client = AI21Client()
+
+    file_response = client.library.files.get(file_in_library)
+    assert set(file_response.labels) == set(DEFAULT_LABELS)
+    sleep(2)
+
+    new_labels = DEFAULT_LABELS + ["new_label"]
+    client.library.files.update(file_in_library, labels=new_labels)
+    file_response = client.library.files.get(file_in_library)
+    assert set(file_response.labels) == set(new_labels)

tests/integration_tests/clients/studio/test_library_answer.py

@@ -0,0 +1,16 @@
+from ai21 import AI21Client
+
+
+def test_library_answer__when_answer_not_in_context__should_return_false(file_in_library: str):
+    client = AI21Client()
+    response = client.library.answer.create(question="Who is Tony Stark?")
+    assert response.answer is None
+    assert not response.answer_in_context
+
+
+def test_library_answer__when_answer_in_context__should_return_true(file_in_library: str):
+    client = AI21Client()
+    response = client.library.answer.create(question="Who was Albert Einstein?")
+    assert response.answer is not None
+    assert response.answer_in_context
+    assert response.sources[0].file_id == file_in_library

tests/integration_tests/clients/studio/test_library_search.py

@@ -0,0 +1,11 @@
+from ai21 import AI21Client
+
+
+def test_library_search__when_search__should_return_relevant_results(file_in_library: str):
+    client = AI21Client()
+    response = client.library.search.create(
+        query="What did Albert Einstein get a Nobel Prize for?", labels=["einstein"]
+    )
+    assert len(response.results) > 0
+    for result in response.results:
+        assert result.file_id == file_in_library