mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-08-29 15:15:03 +02:00
fix migration and add test (#4615)
This commit is contained in:
@@ -127,5 +127,5 @@ def downgrade() -> None:
|
||||
WHERE id = :id
|
||||
"""
|
||||
),
|
||||
{"id": connector_id, "old_config": old_config},
|
||||
{"id": connector_id, "old_config": json.dumps(old_config)},
|
||||
)
|
||||
|
@@ -1,3 +1,5 @@
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
@@ -123,3 +125,196 @@ def test_fix_capitalization_migration() -> None:
|
||||
assert len(results) == 2
|
||||
assert results[0].external_user_group_ids == ["group1", "group2", "group3"]
|
||||
assert results[1].external_user_group_ids == ["upper1", "upper2", "upper3"]
|
||||
|
||||
|
||||
def test_jira_connector_migration() -> None:
|
||||
"""Test that the da42808081e3 migration correctly updates Jira connector configurations"""
|
||||
# Reset the database and run migrations up to the migration before the Jira connector change
|
||||
downgrade_postgres(
|
||||
database="postgres", config_name="alembic", revision="base", clear_data=True
|
||||
)
|
||||
upgrade_postgres(
|
||||
database="postgres",
|
||||
config_name="alembic",
|
||||
# Upgrade it to the migration before the Jira connector change
|
||||
revision="f13db29f3101",
|
||||
)
|
||||
|
||||
# Insert test data with various Jira connector configurations
|
||||
test_data = [
|
||||
{
|
||||
"id": 1,
|
||||
"name": "jira_connector_1",
|
||||
"source": "JIRA",
|
||||
"connector_specific_config": {
|
||||
"jira_project_url": "https://example.atlassian.net/projects/PROJ",
|
||||
"comment_email_blacklist": ["test@example.com"],
|
||||
"batch_size": 100,
|
||||
"labels_to_skip": ["skip-me"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"name": "jira_connector_2",
|
||||
"source": "JIRA",
|
||||
"connector_specific_config": {
|
||||
"jira_project_url": "https://other.atlassian.net/projects/OTHER"
|
||||
},
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"name": "jira_connector_3",
|
||||
"source": "JIRA",
|
||||
"connector_specific_config": {
|
||||
"jira_project_url": "https://example.atlassian.net/projects/TEST",
|
||||
"batch_size": 50,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
# Insert the test data
|
||||
with get_session_context_manager() as db_session:
|
||||
for connector in test_data:
|
||||
db_session.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO connector (
|
||||
id,
|
||||
name,
|
||||
source,
|
||||
connector_specific_config
|
||||
)
|
||||
VALUES (
|
||||
:id,
|
||||
:name,
|
||||
:source,
|
||||
:config
|
||||
)
|
||||
"""
|
||||
),
|
||||
{
|
||||
"id": connector["id"],
|
||||
"name": connector["name"],
|
||||
"source": connector["source"],
|
||||
"config": json.dumps(connector["connector_specific_config"]),
|
||||
},
|
||||
)
|
||||
db_session.commit()
|
||||
|
||||
# Verify the data was inserted correctly
|
||||
with get_session_context_manager() as db_session:
|
||||
results = db_session.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT id, connector_specific_config
|
||||
FROM connector
|
||||
WHERE source = 'JIRA'
|
||||
ORDER BY id
|
||||
"""
|
||||
)
|
||||
).fetchall()
|
||||
|
||||
# Verify initial state
|
||||
assert len(results) == 3
|
||||
assert (
|
||||
results[0].connector_specific_config
|
||||
== test_data[0]["connector_specific_config"]
|
||||
)
|
||||
assert (
|
||||
results[1].connector_specific_config
|
||||
== test_data[1]["connector_specific_config"]
|
||||
)
|
||||
assert (
|
||||
results[2].connector_specific_config
|
||||
== test_data[2]["connector_specific_config"]
|
||||
)
|
||||
|
||||
# Run migrations again to apply the Jira connector change
|
||||
upgrade_postgres(
|
||||
database="postgres", config_name="alembic", revision="da42808081e3"
|
||||
)
|
||||
|
||||
# Verify the upgrade was applied correctly
|
||||
with get_session_context_manager() as db_session:
|
||||
results = db_session.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT id, connector_specific_config
|
||||
FROM connector
|
||||
WHERE source = 'JIRA'
|
||||
ORDER BY id
|
||||
"""
|
||||
)
|
||||
).fetchall()
|
||||
|
||||
# Verify new format
|
||||
assert len(results) == 3
|
||||
|
||||
# First connector - full config
|
||||
config_0 = results[0].connector_specific_config
|
||||
assert config_0["jira_base_url"] == "https://example.atlassian.net"
|
||||
assert config_0["project_key"] == "PROJ"
|
||||
assert config_0["comment_email_blacklist"] == ["test@example.com"]
|
||||
assert config_0["batch_size"] == 100
|
||||
assert config_0["labels_to_skip"] == ["skip-me"]
|
||||
|
||||
# Second connector - minimal config
|
||||
config_1 = results[1].connector_specific_config
|
||||
assert config_1["jira_base_url"] == "https://other.atlassian.net"
|
||||
assert config_1["project_key"] == "OTHER"
|
||||
assert "comment_email_blacklist" not in config_1
|
||||
assert "batch_size" not in config_1
|
||||
assert "labels_to_skip" not in config_1
|
||||
|
||||
# Third connector - partial config
|
||||
config_2 = results[2].connector_specific_config
|
||||
assert config_2["jira_base_url"] == "https://example.atlassian.net"
|
||||
assert config_2["project_key"] == "TEST"
|
||||
assert config_2["batch_size"] == 50
|
||||
assert "comment_email_blacklist" not in config_2
|
||||
assert "labels_to_skip" not in config_2
|
||||
|
||||
# Test downgrade path
|
||||
downgrade_postgres(
|
||||
database="postgres", config_name="alembic", revision="f13db29f3101"
|
||||
)
|
||||
|
||||
# Verify the downgrade was applied correctly
|
||||
with get_session_context_manager() as db_session:
|
||||
results = db_session.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT id, connector_specific_config
|
||||
FROM connector
|
||||
WHERE source = 'JIRA'
|
||||
ORDER BY id
|
||||
"""
|
||||
)
|
||||
).fetchall()
|
||||
|
||||
# Verify reverted to old format
|
||||
assert len(results) == 3
|
||||
|
||||
# First connector - full config
|
||||
config_0 = results[0].connector_specific_config
|
||||
assert (
|
||||
config_0["jira_project_url"]
|
||||
== "https://example.atlassian.net/projects/PROJ"
|
||||
)
|
||||
assert config_0["comment_email_blacklist"] == ["test@example.com"]
|
||||
assert config_0["batch_size"] == 100
|
||||
assert config_0["labels_to_skip"] == ["skip-me"]
|
||||
|
||||
# Second connector - minimal config
|
||||
config_1 = results[1].connector_specific_config
|
||||
assert (
|
||||
config_1["jira_project_url"] == "https://other.atlassian.net/projects/OTHER"
|
||||
)
|
||||
|
||||
# Third connector - partial config
|
||||
config_2 = results[2].connector_specific_config
|
||||
assert (
|
||||
config_2["jira_project_url"]
|
||||
== "https://example.atlassian.net/projects/TEST"
|
||||
)
|
||||
assert config_2["batch_size"] == 50
|
||||
|
Reference in New Issue
Block a user