test

Author: andrewtruong

tests/trace/test_op_return_forms.py

@@ -1,7 +1,7 @@
 import pytest
 
 import weave
-from weave.trace.op_extensions.accumulator import add_accumulator
+from weave.trace.op import _add_accumulator
 from weave.trace.weave_client import get_ref
 from weave.trace_server import trace_server_interface as tsi
 
@@ -110,7 +110,7 @@ def fn():
             size -= 1
             yield size
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     for item in fn():
         pass
@@ -137,7 +137,7 @@ async def fn():
             size -= 1
             yield size
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     async for item in fn():
         pass
@@ -172,7 +172,7 @@ def __next__(self):
     def fn():
         return MyIterator()
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     for item in fn():
         pass
@@ -208,7 +208,7 @@ async def __anext__(self):
     def fn():
         return MyAsyncIterator()
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     async for item in fn():
         pass
@@ -234,7 +234,7 @@ def fn():
             size -= 1
             yield size
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     fn()
 
@@ -260,7 +260,7 @@ async def fn():
             size -= 1
             yield size
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     fn()
 
@@ -294,7 +294,7 @@ def __next__(self):
     def fn():
         return MyIterator()
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     fn()
 
@@ -329,7 +329,7 @@ async def __anext__(self):
     def fn():
         return MyAsyncIterator()
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     fn()
 
@@ -354,7 +354,7 @@ def fn():
             size -= 1
             yield size
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     for item in fn():
         if item == 5:
@@ -382,7 +382,7 @@ async def fn():
             size -= 1
             yield size
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     async for item in fn():
         if item == 5:
@@ -418,7 +418,7 @@ def __next__(self):
     def fn():
         return MyIterator()
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     for item in fn():
         if item == 5:
@@ -455,7 +455,7 @@ async def __anext__(self):
     def fn():
         return MyAsyncIterator()
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     async for item in fn():
         if item == 5:
@@ -484,7 +484,7 @@ def fn():
             if size == 5:
                 raise ValueError("test")
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     try:
         for item in fn():
@@ -517,7 +517,7 @@ async def fn():
             if size == 5:
                 raise ValueError("test")
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     try:
         async for item in fn():
@@ -558,7 +558,7 @@ def __next__(self):
     def fn():
         return MyIterator()
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     try:
         for item in fn():
@@ -600,7 +600,7 @@ async def __anext__(self):
     def fn():
         return MyAsyncIterator()
 
-    add_accumulator(fn, lambda inputs: simple_list_accumulator)
+    _add_accumulator(fn, lambda inputs: simple_list_accumulator)
 
     try:
         async for item in fn():

tests/trace/test_tracing_resilience.py

@@ -14,7 +14,7 @@
 from tests.trace.util import DummyTestException
 from weave.trace.context import call_context
 from weave.trace.context.tests_context import raise_on_captured_errors
-from weave.trace.op_extensions.accumulator import add_accumulator
+from weave.trace.op import _add_accumulator
 
 
 def assert_no_current_call():
@@ -143,7 +143,7 @@ def simple_op():
         def make_accumulator(*args, **kwargs):
             raise DummyTestException("FAILURE!")
 
-        add_accumulator(simple_op, make_accumulator=make_accumulator)
+        _add_accumulator(simple_op, make_accumulator=make_accumulator)
 
         return simple_op()
 
@@ -178,7 +178,7 @@ async def simple_op():
         def make_accumulator(*args, **kwargs):
             raise DummyTestException("FAILURE!")
 
-        add_accumulator(simple_op, make_accumulator=make_accumulator)
+        _add_accumulator(simple_op, make_accumulator=make_accumulator)
 
         return simple_op()
 
@@ -211,7 +211,7 @@ def accumulate(*args, **kwargs):
 
             return accumulate
 
-        add_accumulator(simple_op, make_accumulator=make_accumulator)
+        _add_accumulator(simple_op, make_accumulator=make_accumulator)
 
         return simple_op()
 
@@ -251,7 +251,7 @@ def accumulate(*args, **kwargs):
 
             return accumulate
 
-        add_accumulator(simple_op, make_accumulator=make_accumulator)
+        _add_accumulator(simple_op, make_accumulator=make_accumulator)
 
         return simple_op()
 
@@ -290,7 +290,7 @@ def accumulate(*args, **kwargs):
         def should_accumulate(*args, **kwargs):
             raise DummyTestException("FAILURE!")
 
-        add_accumulator(
+        _add_accumulator(
             simple_op,
             make_accumulator=make_accumulator,
             should_accumulate=should_accumulate,
@@ -335,7 +335,7 @@ def accumulate(*args, **kwargs):
         def should_accumulate(*args, **kwargs):
             raise DummyTestException("FAILURE!")
 
-        add_accumulator(
+        _add_accumulator(
             simple_op,
             make_accumulator=make_accumulator,
             should_accumulate=should_accumulate,
@@ -382,7 +382,7 @@ def accumulate(*args, **kwargs):
         def on_finish_post_processor(*args, **kwargs):
             raise DummyTestException("FAILURE!")
 
-        add_accumulator(
+        _add_accumulator(
             simple_op,
             make_accumulator=make_accumulator,
             on_finish_post_processor=on_finish_post_processor,
@@ -428,7 +428,7 @@ def accumulate(*args, **kwargs):
         def on_finish_post_processor(*args, **kwargs):
             raise DummyTestException("FAILURE!")
 
-        add_accumulator(
+        _add_accumulator(
             simple_op,
             make_accumulator=make_accumulator,
             on_finish_post_processor=on_finish_post_processor,
@@ -467,7 +467,7 @@ def accumulate(*args, **kwargs):
 
             return accumulate
 
-        add_accumulator(simple_op, make_accumulator=make_accumulator)
+        _add_accumulator(simple_op, make_accumulator=make_accumulator)
 
         return simple_op()
 
@@ -497,7 +497,7 @@ def accumulate(*args, **kwargs):
 
             return accumulate
 
-        add_accumulator(simple_op, make_accumulator=make_accumulator)
+        _add_accumulator(simple_op, make_accumulator=make_accumulator)
 
         return simple_op()