Ë
    ðÝjK  ã                   óz   — d Z ddlZddlZddlmZ ddlmZ d„ Z G d„ de«      Ze	dk(  r ej                  «        yy)	zBParallele Anfragen und doppelte Telegram-Nachrichten (Idempotenz).é    N)ÚHandling)ÚGatewayTestCasec                 ó  ‡‡‡— t        j                  | «      Šd g| z  Šˆˆˆfd„}t        | «      D �cg c]  }t        j                  ||f¬«      ‘Œ }}|D ]  }|j	                  «        Œ |D ]  }|j                  d«       Œ ‰S c c}w )Nc                 ó<   •— ‰j                  «         ‰| «      ‰| <   y ©N)Úwait)ÚindexÚbarrierÚresultsÚtargets    €€€ú7/home/aria/aria-v3-dev/python/tests/test_concurrency.pyÚworkerzrun_parallel.<locals>.worker   s   ø€ Ø�‰ŒÙ ›ˆ�Šó    )r   Úargsé<   )Ú	threadingÚBarrierÚrangeÚThreadÚstartÚjoin)Úcountr   r   ÚiÚthreadsÚthreadr
   r   s    `    @@r   Úrun_parallelr   
   sŠ   ú€ Ü×Ñ Ó&€GØˆf�u‰n€Gö'ô DIÈÃ<ÖP¸aŒy×Ñ v°Q°DÖ9ÐP€GÐPØò ˆØ�‰�ðàò ˆØ�‰�B�ðà€Nùò Qs   ³Bc                   ó*   — e Zd Zd„ Zd„ Zd„ Zd„ Zd„ Zy)ÚConcurrencyTestsc                 óÊ  ‡ ‡— ‰ j                  «       Št        dˆˆ fd„«      }‰ j                  t        d„ |D «       «      «       ‰ j	                  t        ‰j                  d   j                  j                  «      d«       ‰ j	                  ‰j                  «       d   d   d   d   d«       ‰ j	                  t        |D �ch c]  }|j                  ’Œ c}«      d«       y c c}w )Né   c                 óL   •— ‰j                  ‰j                  d| z   d«      «      S )Néd   ú#public Frage©ÚhandleÚrequest©r   ÚgatewayÚselfs    €€r   ú<lambda>zBConcurrencyTests.test_parallel_distinct_requests.<locals>.<lambda>   s    ø€ ¨W¯^©^¸D¿L¹LÈÈqÉÐRaÓ<bÓ-c€ r   c              3   óv   K  — | ]1  }|j                   t        j                  u xr |j                  d k(  –— Œ3 y­w)ÚcodexN)Úhandlingr   ÚREPLYÚbackend)Ú.0Úrs     r   ú	<genexpr>zCConcurrencyTests.test_parallel_distinct_requests.<locals>.<genexpr>   s/   è ø€ ÒcÐVW˜AŸJ™J¬(¯.©.Ð8ÒQ¸Q¿Y¹YÈ'Ñ=QÓQÑcùs   ‚79r,   ÚbackendsÚbudgetÚjobs)r(   r   Ú
assertTrueÚallÚassertEqualÚlenÚadaptersÚrunnerÚcallsÚstatusÚjob_id)r)   r   r1   r(   s   `  @r   Útest_parallel_distinct_requestsz0ConcurrencyTests.test_parallel_distinct_requests   s²   ù€ Ø—,‘,“.ˆÜ˜rÔ#cÓdˆØ�‰œÑcÐ[bÔcÓcÔdØ×Ñœ˜W×-Ñ-¨gÑ6×=Ñ=×CÑCÓDÀbÔIØ×Ñ˜Ÿ™Ó)¨*Ñ5°gÑ>¸xÑHÈÑPÐRTÔUØ×Ñœ°Ö8¨1˜aŸh›hÒ8Ó9¸2Õ>ùÒ8s   Ã C c                 ó¶  ‡ ‡— ‰ j                  «       Št        dˆˆ fd„«      }|D �cg c]!  }|j                  t        j                  u sŒ |‘Œ# }}|D �cg c]  }|j
                  dk(  sŒ|‘Œ }}‰ j                  t        |«      t        |«      fd«       ‰ j                  t        ‰j                  d   j                  j                  «      d«       y c c}w c c}w )Né
   c                 óF   •— ‰j                  ‰j                  dd«      «      S )Né   r#   r$   r'   s    €€r   r*   zNConcurrencyTests.test_parallel_duplicate_requests_idempotent.<locals>.<lambda>%   s   ø€ ¨W¯^©^¸D¿L¹LÈÈOÓ<\Ó-]€ r   ÚDUPLICATE_MESSAGE)é   é	   r,   rE   )r(   r   r-   r   r.   Úcoder8   r9   r:   r;   r<   )r)   r   r1   ÚhandledÚdroppedr(   s   `    @r   Ú+test_parallel_duplicate_requests_idempotentz<ConcurrencyTests.test_parallel_duplicate_requests_idempotent#   s®   ù€ Ø—,‘,“.ˆÜ˜rÔ#]Ó^ˆØ%ÖF˜¨¯©´x·~±~Ò)E’1ÐFˆÐFØ%ÖG˜¨¯©Ð3FÓ)F’1ÐGˆÐGØ×Ñœ#˜g›,¬¨G«Ð5°vÔ>Ø×Ñœ˜W×-Ñ-¨gÑ6×=Ñ=×CÑCÓDÀaÕHùò GùÚGs   §!CÁ	CÁCÁ(Cc                 óÊ  — | j                  «       }|j                  | j                  dd«      «      }|j                  | j                  dd«      «      }| j                  |j                  t
        j                  «       | j                  |j                  |j                  ft
        j                  df«       | j                  t        | j                  j                  «      d«       y )Né   ÚhallorD   rE   )r(   r%   r&   r8   r-   r   ÚFORWARD_TO_CLAUDE_CHANNELrG   ÚDROPr9   ÚsinkÚ	delivered)r)   r(   ÚfirstÚseconds       r   Ú$test_sequential_duplicate_idempotentz5ConcurrencyTests.test_sequential_duplicate_idempotent+   sŸ   € Ø—,‘,“.ˆØ—‘˜tŸ|™|¨A¨wÓ7Ó8ˆØ—‘ §¡¨Q°Ó 8Ó9ˆØ×Ñ˜Ÿ™¬×)KÑ)KÔLØ×Ñ˜&Ÿ/™/¨6¯;©;Ð7¼(¿-¹-ÐI\Ð9]Ô^Ø×Ñœ˜TŸY™Y×0Ñ0Ó1°1Õ5r   c                 óV  — | j                  «       j                  | j                  dd«      «       | j                  «       j                  | j                  dd«      «      }| j                  |j                  d«       | j                  t        | j                  j                  «      d«       y )NrF   rM   rD   rE   )r(   r%   r&   r8   rG   r9   rP   rQ   )r)   Úagains     r   Ú%test_duplicate_detected_after_restartz6ConcurrencyTests.test_duplicate_detected_after_restart3   sv   € Ø�‰‹×Ñ˜dŸl™l¨1¨gÓ6Ô7Ø—‘“×%Ñ% d§l¡l°1°gÓ&>Ó?ˆØ×Ñ˜Ÿ™Ð%8Ô9Ø×Ñœ˜TŸY™Y×0Ñ0Ó1°1Õ5r   c                 ó°  ‡ ‡‡— ‰ j                  ddddiii«      Št        dˆˆ fd„«      }|D �cg c]!  }|j                  t        j                  u sŒ |‘Œ# c}Š‰ j                  t        ‰«      d«       ‰ j                  t        ‰j                  d   j                  j                  «      d«       ‰ j                  t        ˆfd„|D «       «      «       y c c}w )Nr4   r,   Ú
daily_jobsé   rA   c                 óL   •— ‰j                  ‰j                  d| z   d«      «      S )NéÈ   z/codex #public Frager$   r'   s    €€r   r*   zCConcurrencyTests.test_parallel_budget_limit_holds.<locals>.<lambda>;   s    ø€ ¨W¯^©^¸D¿L¹LÈÈqÉÐRhÓ<iÓ-j€ r   c              3   ó`   •K  — | ]%  }|‰vsŒ|j                   t        j                  u –— Œ' y ­wr   )r-   r   ÚBLOCKED)r0   r1   Úcodex_repliess     €r   r2   zDConcurrencyTests.test_parallel_budget_limit_holds.<locals>.<genexpr>?   s(   øè ø€ Òf¸qÈqÐXeÒOe˜AŸJ™J¬(×*:Ñ*:Ô:Ñfùs   ƒ	.�!.)r(   r   r-   r   r.   r8   r9   r:   r;   r<   r6   r7   )r)   r   r1   r_   r(   s   `  @@r   Ú test_parallel_budget_limit_holdsz1ConcurrencyTests.test_parallel_budget_limit_holds9   sª   ú€ Ø—,‘, ¨7°\À1Ð4EÐ*FÐGÓHˆÜ˜rÔ#jÓkˆØ$+ÖL˜q¨q¯z©z¼X¿^¹^Ò/KšÒLˆØ×Ñœ˜]Ó+¨QÔ/Ø×Ñœ˜W×-Ñ-¨gÑ6×=Ñ=×CÑCÓDÀaÔHØ�‰œÓfÀGÔfÓfÕgùò Ms   ¯!CÁCN)Ú__name__Ú
__module__Ú__qualname__r?   rJ   rT   rW   r`   © r   r   r   r      s   „ ò?òIò6ò6óhr   r   Ú__main__)Ú__doc__r   ÚunittestÚaria_router.gatewayr   Útests.helpersr   r   r   ra   Úmainrd   r   r   ú<module>rk      sB   ðÙ HÛ Û å (å )òô %h�ô %hðP ˆzÒØ€H‡M�M…Oð r   