
    _j$                        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	m
Z
  e            deddfd	            Z e            deddfd
            Z e            deddfd            Z e            deddfd            Z e            deddfd            Z e            deddfd            Z e            deddfd            Z e            deddfd            ZdS )    )sleep)	ClientAPI)IndexingStatus)skip_if_not_cluster)get_collection_versionwait_for_version_increaseclientreturnNc                    |                                   |                     d          }|                                }t          |t                    sJ |j        dk    sJ |j        dk    sJ |j        dk    sJ |j        dk    sJ dS )z(Test indexing status on empty collectiontest_collectionnamer         ?N)	resetcreate_collectionget_indexing_status
isinstancer   num_indexed_opsnum_unindexed_ops	total_opsop_indexing_progress)r	   
collectionstatuss      c/var/www/FlaskApp/flask-venv/lib/python3.11/site-packages/chromadb/test/api/test_indexing_status.py%test_indexing_status_empty_collectionr      s     LLNNN))/@)AAJ++--Ffn-----!Q&&&&#q((((q    &#------    c                    |                                   |                     d          }d t          d          D             }d t          d          D             }t          | |j                  }|                    ||           |                                }|j        dk    sJ |t          | |j                  k    rt          |t                    sJ |j
        dk    sJ |j        dk    sJ |j        dk    sJ t          | |j        |           t          d	           |                                }t          |t                    sJ |j        dk    sJ |j
        dk    sJ |j        d
k    sJ dS dS )z,Test indexing status after adding embeddingsr   r   c                     g | ]}d | S )id_ .0is     r   
<listcomp>z2test_indexing_status_after_add.<locals>.<listcomp>"   s    
)
)
)999
)
)
)r   ,  c                 r    g | ]4}t          |          t          |d z             t          |dz             g5S       floatr!   s     r   r$   z2test_indexing_status_after_add.<locals>.<listcomp>#   9    MMMQ588U1q5\\5Q<<8MMMr   ids
embeddingsr           <   r   N)r   r   ranger   r   addr   r   r   r   r   r   r   r   r   )r	   r   r.   r/   initial_versionr   final_statuss          r   test_indexing_status_after_addr6      s    LLNNN))/@)AAJ
)
)eCjj
)
)
)CMM%**MMMJ,VZ_EEONNszN222++--Fs""""0IIII&.11111'3....%*****c1111!&*/?KKKb			 "5577,77777+s2222-22220C7777 JI 87r   c                 |   |                                   |                     d          }t          | |j                  }|                    ddgg dg dg           |                                }|j        dk    sJ |t          | |j                  k    rct          |t                    sJ |j	        dk    sJ |j
        d	k    sJ |j        d
k    sJ t          | |j        |           t          d           |                    ddgg dg dg           |                                }|j        dk    sJ dS )z,Test indexing status after upsert operationsr   r   id1id2r   g       @g      @g      @g      @g      @r-   r)   r   r0   r1   id3)g?g @g@g      @g       @g      "@   N)r   r   r   r   upsertr   r   r   r   r   r   r   r   r   )r	   r   r4   r   s       r   !test_indexing_status_after_upsertr@   ;   sf    LLNNN))/@)AAJ,VZ_EEO5%.ooo5WXXX++--Fq    0IIII&.11111'1,,,,%*****c1111!&*/?KKKb			5%.ooo5WXXX++--Fq      r   c                    |                                   |                     d          }t          | |j                  }|                    g dg dg dg dg           |t          | |j                  k    rw|                                }t          |t                    sJ |j        dk    sJ |j	        d	k    sJ |j
        d
k    sJ t          | |j        |           t          d           |                                }|j        dk    sJ |                    ddg           |                                }|j        dk    sJ dS )z,Test indexing status after delete operationsr   r   )r8   r9   r<   r:   r;   r=   r-      r   r0   r1   r8   r9   )r.      N)r   r   r   r   r3   r   r   r   r   r   r   r   r   r   delete)r	   r   r4   r   initial_statusstatus_after_deletes         r   !test_indexing_status_after_deleterG   V   s}    LLNNN))/@)AAJ,VZ_EEONN!!!#OO___oooF    
 0IIII//11&.11111'1,,,,%*****c1111!&*/?KKKb			3355N#q((((5%.))) %88::(A------r   c                 \   |                                   |                     d          }t          | |j                  }|                    dgg dg           |                                }|t          | |j                  k    rct          |t                    sJ |j        dk    sJ |j	        dk    sJ |j
        dk    sJ t          | |j        |           t          d	           |                                }t          |j	        t                    sJ t          |j        t                    sJ t          |j        t                    sJ t          |j
        t                    sJ |j	        dk    sJ |j        dk    sJ |j        dk    sJ d|j
        cxk    rd
k    sn J dS )z5Test that indexing status returns correct field typesfield_types_collectionr   type_test_idr:   r-   r(   r   r0   r1   r   N)r   r   r   r   r3   r   r   r   r   r   r   r   r   intr   r+   )r	   r   r4   r   r5   s        r    test_indexing_status_field_typesrL   v   s    LLNNN))/G)HHJ,VZ_EEONN'___4ENFFF++--F0IIII&.11111'1,,,,%*****c1111!&*/?KKKb			1133Ll2C88888l4c:::::l,c22222l7?????'1,,,,)Q....!Q&&&&,3::::s::::::::r   c                    |                                   |                     d          }t          | |j                   d t	          d          D             }d t	          d          D             }|                    ||           d t	          dd          D             }d	 t	          dd          D             }|                    ||           t          | |j                  }t          d
ddd          t          dddd          t          dd
dd          g}d
}|dk     r^|                                }	|	|v sJ t          d|	           |	j	        }t          | |j        |           t          d           |dk     \dS dS )zITest indexing status with 2000 records based on index version progressionbatch_test_collectionr   c                     g | ]}d | S 	batch_id_r    r!   s     r   r$   z:test_indexing_status_batch_progression.<locals>.<listcomp>   s    222___222r   i  c                 r    g | ]4}t          |          t          |d z             t          |dz             g5S r'   r*   r!   s     r   r$   z:test_indexing_status_batch_progression.<locals>.<listcomp>   s9    PPPqU1XXuQU||U1q5\\:PPPr   r-   c                     g | ]}d | S rP   r    r!   s     r   r$   z:test_indexing_status_batch_progression.<locals>.<listcomp>   s    888___888r   i  c                 r    g | ]4}t          |          t          |d z             t          |dz             g5S r'   r*   r!   s     r   r$   z:test_indexing_status_batch_progression.<locals>.<listcomp>   s9    VVVqU1XXuQU||U1q5\\:VVVr   r   r0   )r   r   r   r   g      ?r   zwitnessed status: r1   N)r   r   r   r   r2   r3   r   r   printr   r   r   )
r	   r   ids_1embeddings_1ids_2embeddings_2current_versionallowed_statusesops_indexedr   s
             r   &test_indexing_status_batch_progressionr]      s    LLNNN))/F)GGJ6:?333 32eDkk222EPPE$KKPPPLNNuN66688eD$&7&7888EVVE$PTDUDUVVVLNNuN666,VZ_EEO 	"!$		
 	
 	
 	 "!$		
 	
 	
 	 !$		
 	
 	
* K


//11)))))"F+++,!&*/?KKKb			 





r   c                 .   |                                   |                     d          }|                     d           	 |                                 J d            # t          $ r}dt          |                                          v skdt          |                                          v sNdt          |                                          v s1dt          |                                          v sJ Y d	}~d	S Y d	}~d	S Y d	}~d	S Y d	}~d	S d	}~ww xY w)
z/Test indexing status on non-existent collectiontemp_collectionr   Fz.Expected exception for non-existent collectionz	not foundzdoes not existzsoft deletedzcollection not foundN)r   r   delete_collectionr   	Exceptionstrlower)r	   r   es      r   test_indexing_status_not_foundre      s6    LLNNN))/@)AAJ
.///	
&&(((FFFFF 
 
 
3q66<<>>))3q66<<>>11Q//%Q7777 *)))))111111//////777777
s   A 
D)BDDc                    |                                   |                     d          }t          | |j                  }d t	          d          D             }d t	          d          D             }|                    ||           t                      }|                                }|                    |j                   |t          | |j                  k    r2|j        dk    sJ t          | |j        |           t          d           |                                }|                    |j                   |j        d	k    sJ t          |          d
k    sJ d|             t          dt          |                      dS )z4Test that progress values vary as indexing completesconcurrent_test_collectionr   c                     g | ]}d | S )concurrent_id_r    r!   s     r   r$   zFtest_indexing_status_concurrent_progress_variation.<locals>.<listcomp>   s!    
4
4
4AA
4
4
4r   r%   c                 r    g | ]4}t          |          t          |d z             t          |dz             g5S r'   r*   r!   s     r   r$   zFtest_indexing_status_concurrent_progress_variation.<locals>.<listcomp>   r,   r   r-   r0   r1   r   r(   z5Expected variation in progress values, but only got: zUnique progress values: N)r   r   r   r   r2   r3   setr   r   r   r   lenrU   sorted)r	   r   r4   r.   r/   progress_valuesr   r5   s           r   2test_indexing_status_concurrent_progress_variationro      s    LLNNN))/K)LLJ,VZ_EEO
4
4s
4
4
4CMM%**MMMJNNszN222eeO ++--F34440IIII*c1111!&*/?KKKb			 1133L9:::,3333 !###QQQ $## 

>VO%<%<
>
>?????r   )timer   chromadb.apir   chromadb.api.typesr   chromadb.test.conftestr   -chromadb.test.utils.wait_for_version_increaser   r   r   r6   r@   rG   rL   r]   re   ro   r    r   r   <module>ru      sI         " " " " " " - - - - - - 6 6 6 6 6 6        .) . . . . . 89 8 8 8 8 8> !i !D ! ! ! !4 .i .D . . . .> ;Y ;4 ; ; ; ;@ .9 . . . . .b 
9 
 
 
 
 
& "@y "@T "@ "@ "@ "@ "@ "@r   