test_rm_kb.py

Overview

test_rm_kb.py is a test suite designed to validate the behavior of the knowledge base removal functionality within the InfiniFlow system. Specifically, it tests the rm_kb function, which appears to handle the deletion of knowledge bases (KBs) from the system, ensuring proper authorization and correct handling of input parameters.

This file uses the pytest framework to define and organize test cases into two main classes:

• TestAuthorization: Tests the authorization mechanisms related to removing knowledge bases.

• TestDatasetsDelete: Tests the deletion functionality under various conditions, including valid and invalid KB IDs.

The tests leverage helper functions rm_kb and list_kbs from the common module and use authorization tokens from libs.auth and configs modules.

Classes and Methods

Class: TestAuthorization

This class contains tests that verify the behavior of the system when removing KBs with invalid or missing authorization.

Method: test_auth_invalid(self, invalid_auth, expected_code, expected_message)

• Purpose:
  Tests that the system correctly rejects attempts to remove a KB when authorization is missing or invalid.

• Parameters:

  • invalid_auth: An authorization object or None to simulate missing or invalid auth.
    Examples used:

    • None (no auth provided)

    • RAGFlowWebApiAuth(INVALID_API_TOKEN) (auth with an explicitly invalid API token)

  • expected_code (int): The expected response code (401 for unauthorized).

  • expected_message (str): The expected error message string.

• Returns:
  None. The method asserts the correctness of the response.

• Behavior:
  Calls rm_kb with the invalid_auth parameter and asserts that the returned response dictionary contains the expected HTTP code (401) and an appropriate unauthorized message.

• Usage Example:

invalid_auth = None
response = rm_kb(invalid_auth)
assert response["code"] == 401
assert "<Unauthorized" in response["message"]

Class: TestDatasetsDelete

This class focuses on testing the deletion functionality of knowledge bases under various scenarios.

Method: test_kb_id(self, WebApiAuth, add_datasets_func)

• Purpose:
  Tests successful deletion of a knowledge base using a valid KB ID.

• Parameters:

  • WebApiAuth: A valid authorization object fixture.

  • add_datasets_func: A fixture that returns a list of KB IDs after adding datasets.

• Returns:
  None. Uses assertions to verify correct deletion.

• Behavior:

  1. Takes the first KB ID from add_datasets_func.

  2. Calls rm_kb with valid authorization and the KB ID payload.

  3. Verifies the removal succeeded (code == 0).

  4. Calls list_kbs to verify the KB count decreased by one.

• Usage Example:

kb_ids = add_datasets_func
payload = {"kb_id": kb_ids[0]}
response = rm_kb(WebApiAuth, payload)
assert response["code"] == 0

response = list_kbs(WebApiAuth)
assert len(response["data"]["kbs"]) == expected_count_after_deletion

Method: test_id_wrong_uuid(self, WebApiAuth)

• Purpose:
  Tests the behavior when an invalid KB ID (wrong UUID format) is provided for deletion.

• Parameters:

  • WebApiAuth: A valid authorization fixture.

  • Uses add_dataset_func fixture to ensure at least one KB exists before test (via @pytest.mark.usefixtures).

• Returns:
  None. Asserts error response correctness.

• Behavior:

  1. Attempts to remove a KB with an invalid UUID string as kb_id.

  2. Asserts that the response code is 109 (custom error code, likely indicating invalid input or authorization failure).

  3. Asserts the error message contains "No authorization."

  4. Calls list_kbs to verify no KBs were deleted (count remains unchanged).

• Usage Example:

invalid_kb_id = "d94a8dc02c9711f0930f7fbc369eab6d"  # malformed UUID
payload = {"kb_id": invalid_kb_id}
response = rm_kb(WebApiAuth, payload)
assert response["code"] == 109
assert "No authorization." in response["message"]

response = list_kbs(WebApiAuth)
assert len(response["data"]["kbs"]) == expected_count_before_deletion

Important Implementation Details and Algorithms

• The tests rely on fixtures such as WebApiAuth, add_datasets_func, and add_dataset_func which are assumed to be defined elsewhere in the test suite. These fixtures handle setup such as authentication and populating the system with test KBs.

• The authorization tests explicitly check for 401 Unauthorized status codes when no or invalid authentication is provided.

• The deletion tests verify both successful deletion by valid ID and graceful failure on invalid input.

• The custom error code 109 indicates a domain-specific error, likely related to authorization or input validation.

• The tests always verify the post-operation state by calling list_kbs to confirm expected KB counts.

Interaction with Other Parts of the System

• rm_kb: The core function under test, imported from common, responsible for removing a knowledge base. Its implementation is external to this file.

• list_kbs: Also imported from common, used to list existing KBs to verify the effect of deletion operations.

• Authorization classes and tokens:

  • RAGFlowWebApiAuth from libs.auth is used to create authorization objects.

  • INVALID_API_TOKEN from configs provides an invalid token for negative testing.

• Fixtures: The test cases depend on pytest fixtures (WebApiAuth, add_datasets_func, add_dataset_func) for setup, indicating integration with a larger test framework and environment.

This file fits into the system as a regression and validation test to ensure that KB removal functionality respects authorization and input validation rules, preventing unauthorized or invalid deletions.

Visual Diagram

classDiagram
    class TestAuthorization {
        +test_auth_invalid(invalid_auth, expected_code, expected_message)
    }

    class TestDatasetsDelete {
        +test_kb_id(WebApiAuth, add_datasets_func)
        +test_id_wrong_uuid(WebApiAuth)
    }

    TestAuthorization --> rm_kb
    TestDatasetsDelete --> rm_kb
    TestDatasetsDelete --> list_kbs

• Explanation:
  The diagram shows the two test classes and their test methods. Both classes call the rm_kb function to attempt KB removal. TestDatasetsDelete also calls list_kbs to verify the state after deletion operations.

Summary

test_rm_kb.py is a concise yet critical test file verifying the correctness and security of the knowledge base removal functionality in InfiniFlow. It ensures that only authorized users can delete KBs, that deletions with valid IDs succeed, and that invalid requests are properly handled. The tests rely on external fixtures and helper functions, making it part of a larger integration testing framework.