github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/clients/python-wrapper/tests/utests/test_import.py (about)

     1  import time
     2  from datetime import timedelta
     3  
     4  import lakefs_sdk
     5  
     6  from lakefs.exceptions import ImportManagerException
     7  from lakefs.import_manager import ImportManager
     8  from tests.utests.common import get_test_client, expect_exception_context
     9  
    10  
    11  class TestImportManager:
    12      def test_import_start_wait(self, monkeypatch):
    13          clt = get_test_client()
    14          mgr = ImportManager("test_repo", "test_branch", client=clt)
    15  
    16          # Wait before start
    17          with expect_exception_context(ImportManagerException):
    18              mgr.wait()
    19  
    20          with monkeypatch.context():
    21              def monkey_import_start(*_, **__):
    22                  return lakefs_sdk.ImportCreationResponse(id="import_id")
    23  
    24              monkeypatch.setattr(lakefs_sdk.ImportApi, "import_start", monkey_import_start)
    25              res = mgr.start()
    26              assert res == "import_id"
    27  
    28              # try again and expect import in progress
    29              with expect_exception_context(ImportManagerException):
    30                  mgr.start()
    31  
    32              requests = 5
    33              status = lakefs_sdk.ImportStatus(completed=False,
    34                                               update_time=time.time(),
    35                                               ingested_objects=500,
    36                                               metarange_id=None,
    37                                               commit=None,
    38                                               error=None)
    39  
    40              def monkey_import_status(*_, **__):
    41                  nonlocal requests, status
    42                  requests -= 1
    43                  if requests == 0:
    44                      status.completed = True
    45                  return status
    46  
    47              monkeypatch.setattr(lakefs_sdk.ImportApi, "import_status", monkey_import_status)
    48              res = mgr.wait(poll_interval=timedelta())
    49              assert res.completed
    50              assert requests == 0
    51  
    52              # try again and expect no error
    53              mgr.wait()