merge in fix with migration (#24314)

Co-authored-by: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Co-authored-by: Ville Brofeldt <ville.brofeldt@apple.com>

tests/integration_tests/key_value/commands/create_test.py

@@ -16,20 +16,23 @@
 # under the License.
 from __future__ import annotations
 
+import json
 import pickle
-from uuid import UUID
 
+import pytest
 from flask.ctx import AppContext
 from flask_appbuilder.security.sqla.models import User
 
 from superset.extensions import db
+from superset.key_value.exceptions import KeyValueCreateFailedError
 from superset.utils.core import override_user
 from tests.integration_tests.key_value.commands.fixtures import (
     admin,
-    ID_KEY,
+    JSON_CODEC,
+    JSON_VALUE,
+    PICKLE_CODEC,
+    PICKLE_VALUE,
     RESOURCE,
-    UUID_KEY,
-    VALUE,
 )
 
 
@@ -38,11 +41,15 @@ def test_create_id_entry(app_context: AppContext, admin: User) -> None:
     from superset.key_value.models import KeyValueEntry
 
     with override_user(admin):
-        key = CreateKeyValueCommand(resource=RESOURCE, value=VALUE).run()
+        key = CreateKeyValueCommand(
+            resource=RESOURCE,
+            value=JSON_VALUE,
+            codec=JSON_CODEC,
+        ).run()
         entry = (
             db.session.query(KeyValueEntry).filter_by(id=key.id).autoflush(False).one()
         )
-        assert pickle.loads(entry.value) == VALUE
+        assert json.loads(entry.value) == JSON_VALUE
         assert entry.created_by_fk == admin.id
         db.session.delete(entry)
         db.session.commit()
@@ -53,11 +60,43 @@ def test_create_uuid_entry(app_context: AppContext, admin: User) -> None:
     from superset.key_value.models import KeyValueEntry
 
     with override_user(admin):
-        key = CreateKeyValueCommand(resource=RESOURCE, value=VALUE).run()
+        key = CreateKeyValueCommand(
+            resource=RESOURCE, value=JSON_VALUE, codec=JSON_CODEC
+        ).run()
     entry = (
         db.session.query(KeyValueEntry).filter_by(uuid=key.uuid).autoflush(False).one()
     )
-    assert pickle.loads(entry.value) == VALUE
+    assert json.loads(entry.value) == JSON_VALUE
     assert entry.created_by_fk == admin.id
     db.session.delete(entry)
     db.session.commit()
+
+
+def test_create_fail_json_entry(app_context: AppContext, admin: User) -> None:
+    from superset.key_value.commands.create import CreateKeyValueCommand
+
+    with pytest.raises(KeyValueCreateFailedError):
+        CreateKeyValueCommand(
+            resource=RESOURCE,
+            value=PICKLE_VALUE,
+            codec=JSON_CODEC,
+        ).run()
+
+
+def test_create_pickle_entry(app_context: AppContext, admin: User) -> None:
+    from superset.key_value.commands.create import CreateKeyValueCommand
+    from superset.key_value.models import KeyValueEntry
+
+    with override_user(admin):
+        key = CreateKeyValueCommand(
+            resource=RESOURCE,
+            value=PICKLE_VALUE,
+            codec=PICKLE_CODEC,
+        ).run()
+        entry = (
+            db.session.query(KeyValueEntry).filter_by(id=key.id).autoflush(False).one()
+        )
+        assert type(pickle.loads(entry.value)) == type(PICKLE_VALUE)
+        assert entry.created_by_fk == admin.id
+        db.session.delete(entry)
+        db.session.commit()