o
    i2                     @   sd  U d dl Z d dlZd dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlm	Z	 d dlm
Z
 d d	lmZ d d
lmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ eeZG dd deZG dd deZejejejejiZ e!ee"f e#d< ejej$ejej%iZ&e!ee"f e#d< 	d@dedej'de!e"e
f dede(f
ddZ)d e j*j+d!e,e
 fd"d#Z-dej'd$ed!efd%d&Z.d'ej/fd e j*j+d(e"d)e"d*e"d+e,e" d,e,e" d-e"d.e(d/e"fd0d1Z0	d@d e j*j+d2e"d3e"d+e,e" d,e,e" d-e"d4e(fd5d6Z1d e j*j+fd7d8Z2d9e
d!e	fd:d;Z3d9e
d!e	fd<d=Z4d e j*j+fd>d?Z5dS )A    N)RetryOutcomes)RetryReason)UserProperty)_get_outcome_from_retry)_get_retry_attempt_string)%_pytest_report_teststatus_return_type)pytest_TestReport)_PYTEST_STATUS)	TestPhase)_get_test_id_from_item)_TestOutcome)get_user_property)TestId)
TestStatus)
get_logger)InternalTestc                   @   s   e Zd ZdZdZdZdS )_ATR_RETRY_OUTCOMESdd_atr_attempt_passeddd_atr_attempt_faileddd_atr_attempt_skippedN)__name__
__module____qualname__ATR_ATTEMPT_PASSEDATR_ATTEMPT_FAILEDATR_ATTEMPT_SKIPPED r   r   ^/home/ubuntu/.local/lib/python3.10/site-packages/ddtrace/contrib/internal/pytest/_atr_utils.pyr      s    r   c                   @   s    e Zd ZdZdZdZdZdZdS )_QUARANTINE_ATR_RETRY_OUTCOMES dd_quarantine_atr_attempt_passed dd_quarantine_atr_attempt_failed!dd_quarantine_atr_attempt_skippeddd_quarantine_atr_final_passeddd_quarantine_atr_final_failedN)r   r   r   r   r   r   ATR_FINAL_PASSEDATR_FINAL_FAILEDr   r   r   r   r      s    r   _FINAL_OUTCOMES_QUARANTINE_FINAL_OUTCOMESFtest_iditemtest_reportstest_outcomeis_quarantinedc                 C   s  | tj}| tj}| tj}|rt}t}	nt}t}	t	|j
|j|j|j
|jd}
|jj|d |rE|jtjkr>|
j|_|jj|d t||
}t| d}t|j|jdd |jD tj||	| |jtjtjftj |j!fg d}|jj|d |jj|d d S )N)PASSEDFAILEDSKIPPEDXFAILXPASSreportfailure_longreprc                 S   s   i | ]}|d qS )   r   ).0kr   r   r   
<dictcomp>\   s    z&atr_handle_retries.<locals>.<dictcomp>)nodeidlocationkeywordswhenlongreproutcomeuser_properties)"getr
   SETUPCALLTEARDOWNr   r'   r   r&   r   r   r   r   ihookpytest_runtest_logreportstatusr   FAILr.   r>   _atr_do_retriesr   	stash_getr   r9   r:   r;   r?   r   RETRY_REASONr   AUTO_TEST_RETRYRETRY_FINAL_OUTCOMEvalue)r(   r)   r*   r+   r,   setup_reportcall_reportteardown_reportretry_outcomesfinal_outcomesoutcomesatr_outcomer=   final_reportr   r   r   atr_handle_retries3   sJ   


rV   terminalreporterreturnc                 C   s   |  tjS N)
getreportsr   r   rW   r   r   r   atr_get_failed_reportsk   s   r\   rS   c                 C   sf   t | }t| t|r,tj|dd}t| ||}t|||j|j|j	 t|st
|}|S )NT)start_immediately)r   r   finishatr_should_retryatr_add_retryr   atr_finish_retryrF   skip_reasonexc_infoatr_get_final_status)r)   rS   r(   	retry_numretry_outcomefinal_statusr   r   r   rH   o   s   



	rH   T
status_keystatus_textreport_outcomeraw_stringsmarkedup_stringscolordelete_reportsretry_reasonc	                    s    fdd|  |D }	|di}
|	r]t|	 d| }|| || jj|fi |
ddi | jd| fi |
ddi |	D ]}| jj| fi |
 d|j }| | qCd S d S )Nc                    s    g | ]}t |tj kr|qS r   )r   r   rJ   )r6   r3   ro   r   r   
<listcomp>   s
    z0_atr_write_report_for_status.<locals>.<listcomp>T bold_)	rZ   lenappend_twmarkup	write_sepupperr9   
write_line)rW   rh   ri   rj   rk   rl   rm   rn   ro   reportsmarkup_kwargstextr3   liner   rp   r   _atr_write_report_for_status   s   

" $r   reports_keyreports_textrs   c           
      C   sh   |  |}|di}|rd|d< |r2t| d| }	||	 || jj|	fi | | j|= d S d S )NTrs   rr   )rZ   ru   rv   rw   rx   stats)
rW   r   r   rk   rl   rm   rs   r|   r}   failed_attempts_textr   r   r   _atr_prepare_attempts_strings   s   
	
r   c              	   C   s  | j tjs	d S | jddddd g }g }t| tjdtj||dd t| tjdtj||d	d g }g }t	| tjd||ddd
 t	| tj
d||d	 t	| tjd||d d|}d|}|r|| jjdddd| | jjddd 7 }|dd| d7 }|| jjdddd7 }|dr|d d }| jddddd |r| jd|| jjt|t|  ddd n	| jddddd | jdddd d S )N=zDatadog Auto Test RetriesT)purplers   failedred)rh   ri   rj   rk   rl   rm   passedgreen)rs   skippedyellowz. z, z (total attempts: )r   ) z[0mrt   z!Datadog Auto Test Retries summaryrr   )	fullwidthr   rs   zNo tests were retried.)r   r@   r   r   ry   r   r.   r	   r-   r   r   r   joinrw   rx   endswithr   ru   )rW   raw_summary_stringsmarkedup_summary_stringsraw_attempt_stringsmarkedup_attempts_stringsraw_summary_stringmarkedup_summary_stringr   r   r   &atr_pytest_terminal_summary_post_yield   s   

	
	

r   r3   c                 C   s   | j tjkrtjddt|  dddiffS | j tjkr,tjddt|  dddiffS | j tjkrBtjd	dt|  d
ddiffS d S )Nrz
ATR RETRY r-   r   TRr.   r   sr/   )r>   r   r   r   r   r   r2   r   r   r   atr_get_teststatus  s    r   c                 C   s   | j tjkrtjddt|  dddiffS | j tjkr,tjddt|  dddiffS | j tjkrBtjddt|  dddiffS | j tjkrRtjd	d
ddiffS | j tjkrbtjddddiffS d S )NqzQUARANTINED RETRY r-   blueTQr.   r/   .z QUARANTINED FINAL STATUS: PASSEDFz QUARANTINED FINAL STATUS: FAILED)r>   r   r   r   r   r   r$   r%   r2   r   r   r   quarantine_atr_get_teststatus)  s4   

r   c                 C   sT   | j tjd  | j tjd  | j tjd  | j tjg  | j tjg  d S rY   )r   popr   r   r   r   r$   r%   r[   r   r   r   -quarantine_pytest_terminal_summary_post_yieldK  s
   r   )F)6_pytestpytest,ddtrace.contrib.internal.pytest._retry_utilsr   r   r   r   r   &ddtrace.contrib.internal.pytest._typesr   r   &ddtrace.contrib.internal.pytest._utilsr	   r
   r   r   r   ddtrace.ext.test_visibility.apir   r   ddtrace.internal.loggerr   $ddtrace.internal.test_visibility.apir   r   logr   r   PASSr-   rG   r.   r&   dictstr__annotations__r$   r%   r'   ItemboolrV   terminalTerminalReporterlistr\   rH   rK   r   r   r   r   r   r   r   r   r   r   <module>   s   
 	

8	
"
a"