From c5a34d77dd0ddd99b2204dc92ee0046dad680566 Mon Sep 17 00:00:00 2001 From: rolinchen Date: Tue, 3 Sep 2024 21:39:25 +0800 Subject: [PATCH] resolve: added test case --- .../tests/apis/web/data_source/conftest.py | 32 ++++++++-- .../apis/web/data_source/test_data_source.py | 58 ++++++++++++++----- 2 files changed, 71 insertions(+), 19 deletions(-) diff --git a/src/bk-user/tests/apis/web/data_source/conftest.py b/src/bk-user/tests/apis/web/data_source/conftest.py index 249a0bd86..1ff1d8b3a 100644 --- a/src/bk-user/tests/apis/web/data_source/conftest.py +++ b/src/bk-user/tests/apis/web/data_source/conftest.py @@ -9,7 +9,9 @@ specific language governing permissions and limitations under the License. """ +import base64 import datetime +import io from typing import Any, Dict, List import pytest @@ -24,7 +26,9 @@ from bkuser.idp_plugins.wecom.plugin import WecomIdpPluginConfig from bkuser.plugins.constants import DataSourcePluginEnum from bkuser.plugins.local.models import LocalDataSourcePluginConfig -from django.test.utils import override_settings +from django.conf import settings +from openpyxl.reader.excel import load_workbook +from openpyxl.workbook import Workbook from tests.test_utils.helpers import generate_random_string @@ -130,7 +134,25 @@ def data_source_sync_tasks(data_source) -> List[DataSourceSyncTask]: return [success_task, failed_task, other_tenant_task] -@pytest.fixture(autouse=True) -def _celery_config(): - with override_settings(CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_EAGER_PROPAGATES=True): - yield +@pytest.fixture() +def data_source_sync_task(data_source) -> DataSourceSyncTask: + return DataSourceSyncTask.objects.create( + data_source=data_source, + status=SyncTaskStatus.PENDING, + has_warning=True, + trigger=SyncTaskTrigger.MANUAL, + extras={"async_run": True, "overwrite": False, "incremental": True}, + ) + + +@pytest.fixture() +def user_workbook() -> Workbook: + return load_workbook(settings.BASE_DIR / "tests/assets/fake_users.xlsx") + + +@pytest.fixture() +def encoded_file(user_workbook): + with io.BytesIO() as buffer: + user_workbook.save(buffer) + content = buffer.getvalue() + return base64.b64encode(content).decode("utf-8") diff --git a/src/bk-user/tests/apis/web/data_source/test_data_source.py b/src/bk-user/tests/apis/web/data_source/test_data_source.py index 03e074e8b..483a08722 100644 --- a/src/bk-user/tests/apis/web/data_source/test_data_source.py +++ b/src/bk-user/tests/apis/web/data_source/test_data_source.py @@ -18,10 +18,13 @@ from bkuser.apps.idp.models import Idp, IdpSensitiveInfo from bkuser.apps.sync.constants import SyncTaskStatus from bkuser.apps.sync.models import DataSourceSyncTask +from bkuser.apps.sync.tasks import sync_data_source +from bkuser.common.cache import Cache, CacheEnum, CacheKeyPrefixEnum from bkuser.plugins.constants import DataSourcePluginEnum from bkuser.plugins.local.constants import PasswordGenerateMethod from django.conf import settings from django.core.files.uploadedfile import SimpleUploadedFile +from django.test.utils import override_settings from django.urls import reverse from rest_framework import status @@ -449,19 +452,46 @@ def test_retrieve_other_tenant_data_source_sync_record(self, api_client, data_so class TestDataSourceImportApi: def test_data_source_import_success(self, api_client, data_source): - url = reverse("data_source.import_from_excel", kwargs={"id": data_source.id}) - with open(settings.BASE_DIR / "tests/assets/fake_users.xlsx", "rb") as excel_file: - uploaded_file = SimpleUploadedFile( - name="fake_users.xlsx", - content=excel_file.read(), - content_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", - ) - response = api_client.post( - url, data={"overwrite": False, "incremental": True, "file": uploaded_file}, format="multipart" + with override_settings(CELERY_TASK_ALWAYS_EAGER=True, CELERY_TASK_EAGER_PROPAGATES=True): + with open(settings.BASE_DIR / "tests/assets/fake_users.xlsx", "rb") as excel_file: + uploaded_file = SimpleUploadedFile( + name="fake_users.xlsx", + content=excel_file.read(), + content_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + ) + resp = api_client.post( + reverse("data_source.import_from_excel", kwargs={"id": data_source.id}), + data={"overwrite": False, "incremental": True, "file": uploaded_file}, + format="multipart", ) + sync_task = DataSourceSyncTask.objects.get(data_source=data_source) + + assert resp.status_code == status.HTTP_200_OK + assert sync_task.status == SyncTaskStatus.SUCCESS + assert DataSourceUser.objects.filter(data_source_id=data_source.id).exists() + assert DataSourceDepartment.objects.filter(data_source_id=data_source.id).exists() + + +class TestDataSourceSyncTask: + def test_sync_data_source_success(self, data_source_sync_task, encoded_file): + cache = Cache(CacheEnum.REDIS, CacheKeyPrefixEnum.DATA_SOURCE_SYNC_RAW_DATA) + task_id = data_source_sync_task.id + task_key = "test_key" + + cache.set(task_key, encoded_file) + plugin_init_extra_kwargs = {"task_key": task_key} + sync_data_source(task_id, plugin_init_extra_kwargs) + + task = DataSourceSyncTask.objects.get(id=task_id) + assert task.status == SyncTaskStatus.SUCCESS.value + + def test_sync_data_source_file_not_found(self, data_source_sync_task): + task_id = data_source_sync_task.id + task_key = "non_existing_key" + + plugin_init_extra_kwargs = {"task_key": task_key} + sync_data_source(task_id, plugin_init_extra_kwargs) - sync_task = DataSourceSyncTask.objects.get(data_source=data_source) - assert response.status_code == status.HTTP_200_OK - assert sync_task.status == SyncTaskStatus.SUCCESS - assert DataSourceUser.objects.filter(data_source_id=data_source.id).exists() - assert DataSourceDepartment.objects.filter(data_source_id=data_source.id).exists() + task = DataSourceSyncTask.objects.get(id=task_id) + assert task.status == SyncTaskStatus.FAILED.value + assert "data source sync task file not found" in task.logs