From bfe8085fa57aca4e70f3dfbb3074638f3a26134f Mon Sep 17 00:00:00 2001
From: Andrew Truong <itsandrewtruong@gmail.com>
Date: Mon, 26 Aug 2024 17:48:22 -0400
Subject: [PATCH] split: test_weave_api into trace and legacy components

---
 weave/legacy/tests/test_weave_api.py      | 28 +----------------------
 weave/tests/trace/test_weave_api_trace.py | 24 +++++++++++++++++++
 2 files changed, 25 insertions(+), 27 deletions(-)
 create mode 100644 weave/tests/trace/test_weave_api_trace.py

diff --git a/weave/legacy/tests/test_weave_api.py b/weave/legacy/tests/test_weave_api.py
index 0e20c9acf37..ace3a42825d 100644
--- a/weave/legacy/tests/test_weave_api.py
+++ b/weave/legacy/tests/test_weave_api.py
@@ -1,11 +1,7 @@
 import os
 
-import pytest
-
 import weave
-import weave.legacy.weave.context_state
-import weave.legacy.weave.wandb_api
-import weave.trace.weave_init
+
 
 
 def test_create_list_rename_delete():
@@ -32,25 +28,3 @@ def test_create_list_rename_delete():
     weave.legacy.weave.ops.delete_artifact(art_node)
     arts = weave.use(weave.legacy.weave.ops.local_artifacts())
     assert len(arts) == 0
-
-
-def test_weave_finish_unsets_client(client):
-    @weave.op
-    def foo():
-        return 1
-
-    weave.trace.client_context.weave_client.set_weave_client_global(None)
-    weave.trace.weave_init._current_inited_client = (
-        weave.trace.weave_init.InitializedClient(client)
-    )
-    weave_client = weave.trace.weave_init._current_inited_client.client
-    assert weave.trace.weave_init._current_inited_client is not None
-
-    foo()
-    assert len(list(weave_client.calls())) == 1
-
-    weave.finish()
-
-    foo()
-    assert len(list(weave_client.calls())) == 1
-    assert weave.trace.weave_init._current_inited_client is None
diff --git a/weave/tests/trace/test_weave_api_trace.py b/weave/tests/trace/test_weave_api_trace.py
new file mode 100644
index 00000000000..4081f180b89
--- /dev/null
+++ b/weave/tests/trace/test_weave_api_trace.py
@@ -0,0 +1,24 @@
+import weave
+import weave.trace.weave_init
+
+
+def test_weave_finish_unsets_client(client):
+    @weave.op
+    def foo():
+        return 1
+
+    weave.trace.client_context.weave_client.set_weave_client_global(None)
+    weave.trace.weave_init._current_inited_client = (
+        weave.trace.weave_init.InitializedClient(client)
+    )
+    weave_client = weave.trace.weave_init._current_inited_client.client
+    assert weave.trace.weave_init._current_inited_client is not None
+
+    foo()
+    assert len(list(weave_client.calls())) == 1
+
+    weave.finish()
+
+    foo()
+    assert len(list(weave_client.calls())) == 1
+    assert weave.trace.weave_init._current_inited_client is None