40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import threading
|
|
from src import ai_client
|
|
|
|
def test_ai_client_tier_isolation():
|
|
original_append = ai_client._append_comms
|
|
captured_logs = []
|
|
def intercepted_append(direction, kind, payload):
|
|
captured_logs.append({
|
|
'thread_name': threading.current_thread().name,
|
|
'source_tier': ai_client.get_current_tier()
|
|
})
|
|
original_append(direction, kind, payload)
|
|
ai_client._append_comms = intercepted_append
|
|
ai_client.clear_comms_log()
|
|
barrier = threading.Barrier(2)
|
|
def run_t1():
|
|
ai_client.set_current_tier('Tier 3')
|
|
barrier.wait()
|
|
ai_client._append_comms('OUT', 'test', {})
|
|
def run_t2():
|
|
ai_client.set_current_tier('Tier 4')
|
|
barrier.wait()
|
|
ai_client._append_comms('OUT', 'test', {})
|
|
t1 = threading.Thread(target=run_t1, name='T1')
|
|
t2 = threading.Thread(target=run_t2, name='T2')
|
|
t1.start()
|
|
t2.start()
|
|
t1.join()
|
|
t2.join()
|
|
try:
|
|
log_t1 = next(log for log in captured_logs if log['thread_name'] == 'T1')
|
|
log_t2 = next(log for log in captured_logs if log['thread_name'] == 'T2')
|
|
assert log_t1['source_tier'] == 'Tier 3'
|
|
assert log_t2['source_tier'] == 'Tier 4'
|
|
finally:
|
|
ai_client._append_comms = original_append
|
|
|
|
if __name__ == "__main__":
|
|
test_ai_client_tier_isolation()
|