From 0c4db8b0f87d8613aac965372fac7fd21c34f9c1 Mon Sep 17 00:00:00 2001
From: Chia-hung Tai <glaxy.tai@gmail.com>
Date: Sat, 19 Nov 2022 11:56:30 +0800
Subject: [PATCH] Set _artifact_uri when mlflow_run is not None. (#1367)

* Set _artifact_uri when mlflow_run is not None.

* Fix black.
---
 qlib/workflow/recorder.py |  1 +
 tests/test_workflow.py    | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)
 create mode 100644 tests/test_workflow.py

diff --git a/qlib/workflow/recorder.py b/qlib/workflow/recorder.py
index 9d82bf0a4..33903f7ed 100644
--- a/qlib/workflow/recorder.py
+++ b/qlib/workflow/recorder.py
@@ -279,6 +279,7 @@ class MLflowRecorder(Recorder):
                 if mlflow_run.info.end_time is not None
                 else None
             )
+            self._artifact_uri = mlflow_run.info.artifact_uri
         self.async_log = None
 
     def __repr__(self):
diff --git a/tests/test_workflow.py b/tests/test_workflow.py
new file mode 100644
index 000000000..129abc0fb
--- /dev/null
+++ b/tests/test_workflow.py
@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+import unittest
+from pathlib import Path
+import shutil
+
+from qlib.workflow import R
+from qlib.tests import TestAutoData
+
+
+class WorkflowTest(TestAutoData):
+    TMP_PATH = Path("./.mlruns_tmp/")
+
+    def tearDown(self) -> None:
+        if self.TMP_PATH.exists():
+            shutil.rmtree(self.TMP_PATH)
+
+    def test_get_local_dir(self):
+        """ """
+        with R.start(uri=str(self.TMP_PATH)):
+            pass
+
+        with R.uri_context(uri=str(self.TMP_PATH)):
+            resume_recorder = R.get_recorder()
+            resume_recorder.get_local_dir()
+
+
+if __name__ == "__main__":
+    unittest.main()