From 8174c694a8734ef73c65674f035675a154b407f9 Mon Sep 17 00:00:00 2001
From: Aaron Kaplan <33381112+esoteric-ephemera@users.noreply.github.com>
Date: Tue, 20 Aug 2024 02:04:07 -0700
Subject: [PATCH] Fix MPID assignment in electrode workflow (#956)

* fix mpid in electrode

* fix mpid in electrode

* ensure unique mpid

* precommit

* change mpid assignment logic to assign ULIDs on the fly as needed

* precommit

* remove missing test fixture call
---
 src/atomate2/common/jobs/electrode.py | 16 ++++++++++++++--
 tests/vasp/flows/test_electrode.py    | 21 +--------------------
 2 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/src/atomate2/common/jobs/electrode.py b/src/atomate2/common/jobs/electrode.py
index 1d3786c077..4ef3a29a02 100644
--- a/src/atomate2/common/jobs/electrode.py
+++ b/src/atomate2/common/jobs/electrode.py
@@ -6,10 +6,12 @@
 from typing import TYPE_CHECKING, Callable, NamedTuple
 
 from emmet.core.electrode import InsertionElectrodeDoc
+from emmet.core.mpid import MPID
 from emmet.core.structure_group import StructureGroupDoc
 from jobflow import Flow, Maker, Response, job
 from pymatgen.analysis.defects.generators import ChargeInterstitialGenerator
 from pymatgen.entries.computed_entries import ComputedStructureEntry
+from ulid import ULID
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -154,13 +156,23 @@ def get_computed_entries(
         return multi
     # keep the [1] for now, if jobflow supports NamedTuple, we can do this much cleaner
     s_ = RelaxJobSummary._make(single)
-    s_.entry.entry_id = s_.uuid
+
+    # Ensure that the entry_id is an acceptable MPID
+    try:
+        entry_id = MPID(s_.uuid)
+    except ValueError:
+        entry_id = ULID()
+    s_.entry.entry_id = str(entry_id)
+
+    # Store UUID for provenance
+    s_.entry.data["UUID"] = s_.uuid
+
     ent = ComputedStructureEntry(
         structure=s_.structure,
         energy=s_.entry.energy,
         parameters=s_.entry.parameters,
         data=s_.entry.data,
-        entry_id=s_.uuid,
+        entry_id=s_.entry.entry_id,
     )
     return [*multi, ent]
 
diff --git a/tests/vasp/flows/test_electrode.py b/tests/vasp/flows/test_electrode.py
index 0ce47c579b..79ab6738b2 100644
--- a/tests/vasp/flows/test_electrode.py
+++ b/tests/vasp/flows/test_electrode.py
@@ -1,25 +1,7 @@
 from __future__ import annotations
 
-from unittest import mock
 
-import pytest
-from jobflow.settings import JobflowSettings
-
-
-@pytest.fixture()
-def mock_jobflow_settings(memory_jobstore):
-    """Set the UID_TYPE to "ulid" to make sure the documents can be sorted.
-
-    See: https://github.com/materialsproject/jobflow/issues/519#issuecomment-1906850096
-    """
-
-    settings = JobflowSettings(JOB_STORE=memory_jobstore, UID_TYPE="ulid")
-
-    with mock.patch("jobflow.SETTINGS", settings):
-        yield
-
-
-def test_electrode_makers(mock_vasp, clean_dir, test_dir, mock_jobflow_settings):
+def test_electrode_makers(mock_vasp, clean_dir, test_dir):
     from emmet.core.electrode import InsertionElectrodeDoc
     from jobflow import OutputReference, run_locally
     from monty.serialization import loadfn
@@ -32,7 +14,6 @@ def test_electrode_makers(mock_vasp, clean_dir, test_dir, mock_jobflow_settings)
         update_user_incar_settings,
         update_user_kpoints_settings,
     )
-    # mock the default setting
 
     # mapping from job name to directory containing test files
     ref_paths = {