""" Tests for RAGFlow schema conversion. This module tests: - RAGFlowSchemaConverter: Analyzes ES mappings and generates OB column definitions - RAGFlowDataConverter: Converts ES documents to OceanBase row format - Vector field pattern matching - Schema constants """ import json from es_ob_migration.schema import ( RAGFlowSchemaConverter, RAGFlowDataConverter, RAGFLOW_COLUMNS, ARRAY_COLUMNS, JSON_COLUMNS, VECTOR_FIELD_PATTERN, FTS_COLUMNS_ORIGIN, FTS_COLUMNS_TKS, ) class TestRAGFlowSchemaConverter: """Test RAGFlowSchemaConverter class.""" def test_analyze_ragflow_mapping(self): """Test analyzing a RAGFlow ES mapping.""" converter = RAGFlowSchemaConverter() # Simulate a RAGFlow ES mapping es_mapping = { "properties": { "id": {"type": "keyword"}, "kb_id": {"type": "keyword"}, "doc_id": {"type": "keyword"}, "docnm_kwd": {"type": "keyword"}, "content_with_weight": {"type": "text"}, "content_ltks": {"type": "text"}, "available_int": {"type": "integer"}, "important_kwd": {"type": "keyword"}, "q_768_vec": {"type": "dense_vector", "dims": 768}, } } analysis = converter.analyze_es_mapping(es_mapping) # Check known fields assert "id" in analysis["known_fields"] assert "kb_id" in analysis["known_fields"] assert "content_with_weight" in analysis["known_fields"] # Check vector fields assert len(analysis["vector_fields"]) == 1 assert analysis["vector_fields"][0]["name"] == "q_768_vec" assert analysis["vector_fields"][0]["dimension"] == 768 def test_detect_vector_size(self): """Test automatic vector size detection.""" converter = RAGFlowSchemaConverter() es_mapping = { "properties": { "q_1536_vec": {"type": "dense_vector", "dims": 1536}, } } converter.analyze_es_mapping(es_mapping) assert converter.detected_vector_size == 1536 def test_unknown_fields(self): """Test that unknown fields are properly identified.""" converter = RAGFlowSchemaConverter() es_mapping = { "properties": { "id": {"type": "keyword"}, "custom_field": {"type": "text"}, "another_field": {"type": "integer"}, } } analysis = converter.analyze_es_mapping(es_mapping) assert "custom_field" in analysis["unknown_fields"] assert "another_field" in analysis["unknown_fields"] def test_get_column_definitions(self): """Test getting RAGFlow column definitions.""" converter = RAGFlowSchemaConverter() # First analyze to detect vector fields es_mapping = { "properties": { "q_768_vec": {"type": "dense_vector", "dims": 768}, } } converter.analyze_es_mapping(es_mapping) columns = converter.get_column_definitions() # Check that all RAGFlow columns are present column_names = [c["name"] for c in columns] for col_name in RAGFLOW_COLUMNS: assert col_name in column_names, f"Missing column: {col_name}" # Check vector column is added assert "q_768_vec" in column_names class TestRAGFlowDataConverter: """Test RAGFlowDataConverter class.""" def test_convert_basic_document(self): """Test converting a basic RAGFlow document.""" converter = RAGFlowDataConverter() es_doc = { "_id": "test-id-123", "_source": { "id": "test-id-123", "kb_id": "kb-001", "doc_id": "doc-001", "docnm_kwd": "test_document.pdf", "content_with_weight": "This is test content", "available_int": 1, } } row = converter.convert_document(es_doc) assert row["id"] == "test-id-123" assert row["kb_id"] == "kb-001" assert row["doc_id"] == "doc-001" assert row["docnm_kwd"] == "test_document.pdf" assert row["content_with_weight"] == "This is test content" assert row["available_int"] == 1 def test_convert_with_vector(self): """Test converting document with vector embedding.""" converter = RAGFlowDataConverter() embedding = [0.1] * 768 es_doc = { "_id": "vec-doc-001", "_source": { "id": "vec-doc-001", "kb_id": "kb-001", "q_768_vec": embedding, } } row = converter.convert_document(es_doc) assert row["id"] == "vec-doc-001" assert row["q_768_vec"] == embedding assert "q_768_vec" in converter.vector_fields def test_convert_array_fields(self): """Test converting array fields.""" converter = RAGFlowDataConverter() es_doc = { "_id": "array-doc", "_source": { "id": "array-doc", "kb_id": "kb-001", "important_kwd": ["keyword1", "keyword2", "keyword3"], "question_kwd": ["What is this?", "How does it work?"], "tag_kwd": ["tag1", "tag2"], } } row = converter.convert_document(es_doc) # Array fields should be JSON strings assert isinstance(row["important_kwd"], str) parsed = json.loads(row["important_kwd"]) assert parsed == ["keyword1", "keyword2", "keyword3"] def test_convert_json_fields(self): """Test converting JSON fields.""" converter = RAGFlowDataConverter() es_doc = { "_id": "json-doc", "_source": { "id": "json-doc", "kb_id": "kb-001", "tag_feas": {"tag1": 0.8, "tag2": 0.5}, "metadata": {"author": "John", "date": "2024-01-01"}, } } row = converter.convert_document(es_doc) # JSON fields should be JSON strings assert isinstance(row["tag_feas"], str) assert isinstance(row["metadata"], str) tag_feas = json.loads(row["tag_feas"]) assert tag_feas == {"tag1": 0.8, "tag2": 0.5} def test_convert_unknown_fields_to_extra(self): """Test that unknown fields are stored in 'extra'.""" converter = RAGFlowDataConverter() es_doc = { "_id": "extra-doc", "_source": { "id": "extra-doc", "kb_id": "kb-001", "custom_field": "custom_value", "another_custom": 123, } } row = converter.convert_document(es_doc) assert "extra" in row extra = json.loads(row["extra"]) assert extra["custom_field"] == "custom_value" assert extra["another_custom"] == 123 def test_convert_kb_id_list(self): """Test converting kb_id when it's a list (ES format).""" converter = RAGFlowDataConverter() es_doc = { "_id": "kb-list-doc", "_source": { "id": "kb-list-doc", "kb_id": ["kb-001", "kb-002"], # Some ES docs have list } } row = converter.convert_document(es_doc) # Should take first element assert row["kb_id"] == "kb-001" def test_convert_content_with_weight_dict(self): """Test converting content_with_weight when it's a dict.""" converter = RAGFlowDataConverter() es_doc = { "_id": "content-dict-doc", "_source": { "id": "content-dict-doc", "kb_id": "kb-001", "content_with_weight": { "text": "Some content", "weight": 1.0, }, } } row = converter.convert_document(es_doc) # Dict should be JSON serialized assert isinstance(row["content_with_weight"], str) parsed = json.loads(row["content_with_weight"]) assert parsed["text"] == "Some content" def test_convert_batch(self): """Test batch conversion.""" converter = RAGFlowDataConverter() es_docs = [ {"_id": f"doc-{i}", "_source": {"id": f"doc-{i}", "kb_id": "kb-001"}} for i in range(5) ] rows = converter.convert_batch(es_docs) assert len(rows) == 5 for i, row in enumerate(rows): assert row["id"] == f"doc-{i}" class TestVectorFieldPattern: """Test vector field pattern matching.""" def test_valid_patterns(self): """Test valid vector field patterns.""" valid_names = [ "q_768_vec", "q_1024_vec", "q_1536_vec", "q_3072_vec", ] for name in valid_names: match = VECTOR_FIELD_PATTERN.match(name) assert match is not None, f"Should match: {name}" def test_invalid_patterns(self): """Test invalid vector field patterns.""" invalid_names = [ "q_vec", "768_vec", "q_768", "vector_768", "content_with_weight", ] for name in invalid_names: match = VECTOR_FIELD_PATTERN.match(name) assert match is None, f"Should not match: {name}" def test_extract_dimension(self): """Test extracting dimension from pattern.""" match = VECTOR_FIELD_PATTERN.match("q_1536_vec") assert match is not None assert int(match.group("vector_size")) == 1536 class TestConstants: """Test schema constants.""" def test_array_columns(self): """Test ARRAY_COLUMNS list.""" expected = [ "important_kwd", "question_kwd", "tag_kwd", "source_id", "entities_kwd", "position_int", "page_num_int", "top_int" ] for col in expected: assert col in ARRAY_COLUMNS, f"Missing array column: {col}" def test_json_columns(self): """Test JSON_COLUMNS list.""" expected = ["tag_feas", "metadata", "extra"] for col in expected: assert col in JSON_COLUMNS, f"Missing JSON column: {col}" def test_ragflow_columns_completeness(self): """Test that RAGFLOW_COLUMNS has all required fields.""" required_fields = [ "id", "kb_id", "doc_id", "content_with_weight", "available_int", "metadata", "extra", ] for field in required_fields: assert field in RAGFLOW_COLUMNS, f"Missing required field: {field}" def test_fts_columns(self): """Test fulltext search column lists.""" assert "content_with_weight" in FTS_COLUMNS_ORIGIN assert "content_ltks" in FTS_COLUMNS_TKS def test_ragflow_columns_types(self): """Test column type definitions.""" # Primary key assert RAGFLOW_COLUMNS["id"]["is_primary"] is True assert RAGFLOW_COLUMNS["id"]["nullable"] is False # Indexed columns assert RAGFLOW_COLUMNS["kb_id"]["index"] is True assert RAGFLOW_COLUMNS["doc_id"]["index"] is True # Array columns assert RAGFLOW_COLUMNS["important_kwd"]["is_array"] is True assert RAGFLOW_COLUMNS["question_kwd"]["is_array"] is True # JSON columns assert RAGFLOW_COLUMNS["metadata"]["is_json"] is True assert RAGFLOW_COLUMNS["extra"]["is_json"] is True class TestRAGFlowSchemaConverterEdgeCases: """Test edge cases for RAGFlowSchemaConverter.""" def test_empty_mapping(self): """Test analyzing empty mapping.""" converter = RAGFlowSchemaConverter() analysis = converter.analyze_es_mapping({}) assert analysis["known_fields"] == [] assert analysis["vector_fields"] == [] assert analysis["unknown_fields"] == [] def test_mapping_without_properties(self): """Test mapping without properties key.""" converter = RAGFlowSchemaConverter() analysis = converter.analyze_es_mapping({"some_other_key": {}}) assert analysis["known_fields"] == [] def test_multiple_vector_fields(self): """Test detecting multiple vector fields.""" converter = RAGFlowSchemaConverter() es_mapping = { "properties": { "q_768_vec": {"type": "dense_vector", "dims": 768}, "q_1024_vec": {"type": "dense_vector", "dims": 1024}, } } analysis = converter.analyze_es_mapping(es_mapping) assert len(analysis["vector_fields"]) == 2 # First detected should be set assert converter.detected_vector_size in [768, 1024] def test_get_column_definitions_without_analysis(self): """Test getting columns without prior analysis.""" converter = RAGFlowSchemaConverter() columns = converter.get_column_definitions() # Should have all RAGFlow columns but no vector columns column_names = [c["name"] for c in columns] assert "id" in column_names assert "kb_id" in column_names def test_get_vector_fields(self): """Test getting vector fields.""" converter = RAGFlowSchemaConverter() es_mapping = { "properties": { "q_1536_vec": {"type": "dense_vector", "dims": 1536}, } } converter.analyze_es_mapping(es_mapping) vec_fields = converter.get_vector_fields() assert len(vec_fields) == 1 assert vec_fields[0]["name"] == "q_1536_vec" assert vec_fields[0]["dimension"] == 1536 class TestRAGFlowDataConverterEdgeCases: """Test edge cases for RAGFlowDataConverter.""" def test_convert_empty_document(self): """Test converting empty document.""" converter = RAGFlowDataConverter() es_doc = {"_id": "empty_doc", "_source": {}} row = converter.convert_document(es_doc) assert row["id"] == "empty_doc" def test_convert_document_without_source(self): """Test converting document without _source.""" converter = RAGFlowDataConverter() es_doc = {"_id": "no_source", "id": "no_source", "kb_id": "kb_001"} row = converter.convert_document(es_doc) assert row["id"] == "no_source" assert row["kb_id"] == "kb_001" def test_convert_boolean_to_integer(self): """Test converting boolean to integer.""" converter = RAGFlowDataConverter() es_doc = { "_id": "bool_doc", "_source": { "id": "bool_doc", "kb_id": "kb_001", "available_int": True, } } row = converter.convert_document(es_doc) assert row["available_int"] == 1 def test_convert_invalid_integer(self): """Test converting invalid integer value.""" converter = RAGFlowDataConverter() es_doc = { "_id": "invalid_int", "_source": { "id": "invalid_int", "kb_id": "kb_001", "available_int": "not_a_number", } } row = converter.convert_document(es_doc) assert row["available_int"] is None def test_convert_float_field(self): """Test converting float fields.""" converter = RAGFlowDataConverter() es_doc = { "_id": "float_doc", "_source": { "id": "float_doc", "kb_id": "kb_001", "weight_flt": 0.85, "rank_flt": "0.95", # String that should become float } } row = converter.convert_document(es_doc) assert row["weight_flt"] == 0.85 assert row["rank_flt"] == 0.95 def test_convert_array_with_special_characters(self): """Test converting array with special characters.""" converter = RAGFlowDataConverter() es_doc = { "_id": "special_array", "_source": { "id": "special_array", "kb_id": "kb_001", "important_kwd": ["key\nwith\nnewlines", "key\twith\ttabs"], } } row = converter.convert_document(es_doc) # Should be JSON string with escaped characters assert isinstance(row["important_kwd"], str) parsed = json.loads(row["important_kwd"]) assert len(parsed) == 2 def test_convert_already_json_array(self): """Test converting already JSON-encoded array.""" converter = RAGFlowDataConverter() es_doc = { "_id": "json_array", "_source": { "id": "json_array", "kb_id": "kb_001", "important_kwd": '["already", "json"]', } } row = converter.convert_document(es_doc) assert row["important_kwd"] == '["already", "json"]' def test_convert_single_value_to_array(self): """Test converting single value to array.""" converter = RAGFlowDataConverter() es_doc = { "_id": "single_to_array", "_source": { "id": "single_to_array", "kb_id": "kb_001", "important_kwd": "single_keyword", } } row = converter.convert_document(es_doc) parsed = json.loads(row["important_kwd"]) assert parsed == ["single_keyword"] def test_detect_vector_fields_from_document(self): """Test detecting vector fields from document.""" converter = RAGFlowDataConverter() doc = { "q_768_vec": [0.1] * 768, "q_1024_vec": [0.2] * 1024, } converter.detect_vector_fields(doc) assert "q_768_vec" in converter.vector_fields assert "q_1024_vec" in converter.vector_fields def test_convert_with_default_values(self): """Test conversion uses default values.""" converter = RAGFlowDataConverter() es_doc = { "_id": "default_test", "_source": { "id": "default_test", "kb_id": "kb_001", # available_int not provided, should get default } } row = converter.convert_document(es_doc) # available_int has default of 1 assert row.get("available_int") == 1 def test_convert_list_content(self): """Test converting list content to JSON.""" converter = RAGFlowDataConverter() es_doc = { "_id": "list_content", "_source": { "id": "list_content", "kb_id": "kb_001", "content_with_weight": ["part1", "part2", "part3"], } } row = converter.convert_document(es_doc) assert isinstance(row["content_with_weight"], str) parsed = json.loads(row["content_with_weight"]) assert parsed == ["part1", "part2", "part3"] def test_convert_batch_empty(self): """Test batch conversion with empty list.""" converter = RAGFlowDataConverter() rows = converter.convert_batch([]) assert rows == [] def test_existing_extra_field_merged(self): """Test that existing extra field is merged with unknown fields.""" converter = RAGFlowDataConverter() es_doc = { "_id": "merge_extra", "_source": { "id": "merge_extra", "kb_id": "kb_001", "extra": {"existing_key": "existing_value"}, "custom_field": "custom_value", } } row = converter.convert_document(es_doc) # extra should contain both existing and new fields extra = json.loads(row["extra"]) assert "custom_field" in extra