use hard coded time intervals for provisioning

.github/workflows/CI.yml

@@ -31,7 +31,7 @@ jobs:
     needs: run-linters
     strategy:
       matrix:
-        python-version: [3.9, '3.10']
+        python-version: [3.9, '3.10', 3.11]
 
     steps:
     - uses: actions/checkout@v3

context_engine/knoweldge_base/knowledge_base.py

@@ -24,6 +24,9 @@
 
 
 INDEX_NAME_PREFIX = "context-engine-"
+TIMEOUT_INDEX_CREATE = 300
+TIMEOUT_INDEX_PROVISION = 30
+INDEX_PROVISION_TIME_INTERVAL = 3
 
 
 class KnowledgeBase(BaseKnowledgeBase):
@@ -94,22 +97,10 @@ def create_with_new_index(cls,
                               default_top_k: int = 10,
                               indexed_fields: Optional[List[str]] = None,
                               dimension: Optional[int] = None,
-                              timeout: int = 300,
-                              time_interval: int = 5,
                               create_index_params: Optional[dict] = None
                               ) -> 'KnowledgeBase':
 
         # validate inputs
-
-        if timeout < 0:
-            raise ValueError(
-                f"timeout must be a on-negative integer. Got {timeout}"
-            )
-        if time_interval < 0:
-            raise ValueError(
-                f"time_interval must be a on-negative integer. Got {time_interval}"
-            )
-
         if indexed_fields is None:
             indexed_fields = ['document_id']
         elif "document_id" not in indexed_fields:
@@ -147,6 +138,7 @@ def create_with_new_index(cls,
                          metadata_config={
                              'indexed': indexed_fields
                          },
+                         timeout=TIMEOUT_INDEX_CREATE,
                          **create_index_params)
         except Exception as e:
             raise RuntimeError(
@@ -155,9 +147,7 @@ def create_with_new_index(cls,
             ) from e
 
         # wait for index to be provisioned
-        cls._wait_for_index_provision(full_index_name=full_index_name,
-                                      timeout=timeout,
-                                      time_interval=time_interval)
+        cls._wait_for_index_provision(full_index_name=full_index_name)
 
         # initialize KnowledgeBase
         return cls(index_name=index_name,
@@ -168,9 +158,7 @@ def create_with_new_index(cls,
 
     @classmethod
     def _wait_for_index_provision(cls,
-                                  full_index_name: str,
-                                  timeout: int,
-                                  time_interval: int):
+                                  full_index_name: str):
         start_time = time.time()
         while True:
             try:
@@ -181,13 +169,13 @@ def _wait_for_index_provision(cls,
                 pass
 
             time_passed = time.time() - start_time
-            if time_passed > timeout:
+            if time_passed > TIMEOUT_INDEX_PROVISION:
                 raise RuntimeError(
                     f"Index {full_index_name} failed to provision "
                     f"for {time_passed} seconds."
                     f"Please try creating KnowledgeBase again in a few minutes."
                 )
-            time.sleep(time_interval)
+            time.sleep(INDEX_PROVISION_TIME_INTERVAL)
 
     @staticmethod
     def _get_full_index_name(index_name: str) -> str:

pyproject.toml

@@ -30,7 +30,7 @@ flake8-pyproject = "^1.2.3"
 asyncio = "^3.4.3"
 pytest-asyncio = "^0.14.0"
 pytest-mock = "^3.6.1"
-pytest-order = "^1.1.0"
+
 
 [build-system]
 requires = ["poetry-core"]

tests/system/knowledge_base/test_knowledge_base.py

@@ -16,6 +16,8 @@
 
 load_dotenv()
 
+PINECONE_API_KEY_ENV_VAR = "PINECONE_API_KEY"
+
 
 @pytest.fixture(scope="module")
 def index_name(testrun_uid):
@@ -84,9 +86,6 @@ def assert_ids_not_in_index(knowledge_base, ids):
 def teardown_knowledge_base(index_full_name, knowledge_base):
     yield
 
-    # some tests change the env vars, so we reload them here
-    # load_dotenv(override=True)
-
     pinecone.init()
     if index_full_name in pinecone.list_indexes():
         pinecone.delete_index(index_full_name)
@@ -263,31 +262,19 @@ def test_delete_index_for_non_existing(knowledge_base):
     assert "index was deleted." in str(e.value)
 
 
-@pytest.mark.parametrize("timeout, time_interval, indexed_fields, error_msg", [
-    (-1, 1, ["id"], "timeout must be a on-negative integer"),
-    (1, -1, ["id"], "time_interval must be a on-negative integer"),
-    (1, 1, ["id", "text", "metadata"],
-     "The 'text' field cannot be used for metadata filtering"),
-])
-def test_create_with_new_index_invalid_params(index_name,
-                                              chunker,
-                                              encoder,
-                                              timeout,
-                                              time_interval,
-                                              indexed_fields,
-                                              error_msg):
+def test_create_with_text_in_indexed_field_raise(index_name,
+                                                 chunker,
+                                                 encoder):
     with pytest.raises(ValueError) as e:
         KnowledgeBase.create_with_new_index(index_name=index_name,
                                             encoder=encoder,
                                             chunker=chunker,
-                                            timeout=timeout,
-                                            time_interval=time_interval,
-                                            indexed_fields=indexed_fields)
+                                            indexed_fields=["id", "text", "metadata"])
 
-    assert error_msg in str(e.value)
+    assert "The 'text' field cannot be used for metadata filtering" in str(e.value)
 
 
-def test_create_with_new_index_ecoder_dimension_none(index_name, chunker):
+def test_create_with_new_index_encoder_dimension_none(index_name, chunker):
     encoder = StubRecordEncoder(StubDenseEncoder(dimension=3))
     encoder._dense_encoder.dimension = None
     with pytest.raises(ValueError) as e:
@@ -298,9 +285,19 @@ def test_create_with_new_index_ecoder_dimension_none(index_name, chunker):
     assert "Could not infer dimension from encoder" in str(e.value)
 
 
-@pytest.mark.skip(reason="This test is flaky")
-def test_create_bad_credentials(index_name, chunker, encoder):
-    os.environ["PINECONE_API_KEY"] = "bad-key"
+@pytest.fixture
+def set_bad_credentials():
+    original_api_key = os.environ.get(PINECONE_API_KEY_ENV_VAR)
+
+    os.environ[PINECONE_API_KEY_ENV_VAR] = "bad-key"
+
+    yield
+
+    # Restore the original API key after test execution
+    os.environ[PINECONE_API_KEY_ENV_VAR] = original_api_key
+
+
+def test_create_bad_credentials(set_bad_credentials, index_name, chunker, encoder):
     with pytest.raises(RuntimeError) as e:
         KnowledgeBase.create_with_new_index(index_name=index_name,
                                             encoder=encoder,
@@ -309,9 +306,7 @@ def test_create_bad_credentials(index_name, chunker, encoder):
     assert "Please check your credentials" in str(e.value)
 
 
-@pytest.mark.skip(reason="This test is flaky")
-def test_init_bad_credentials(index_name, chunker, encoder):
-    os.environ["PINECONE_API_KEY"] = "bad-key"
+def test_init_bad_credentials(set_bad_credentials, index_name, chunker, encoder):
     with pytest.raises(RuntimeError) as e:
         KnowledgeBase(index_name=index_name,
                       encoder=encoder,