split: test_weave_api into trace and legacy components

wandb · Aug 26, 2024 · bfe8085 · bfe8085
1 parent ae73ae5
commit bfe8085
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 27 deletions.
diff --git 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
@@ -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