ó
    ]`0j†;  ã                   óº   • S r SSKrSSKJs  Jr  SSKrSSKJr  SSK	J
r
  SSKJrJrJr  SSKrSSKJrJrJr  SSKJr   " S S	5      r " S
 S5      r " S S5      rg)z$Tests for JSON import functionality.é    N)Údatetime)ÚPath)Ú	MagicMockÚ	mock_openÚpatch)Úget_import_statsÚimport_directoryÚimport_json_file)ÚDocumentc                   óº  • \ rS rSrSr\R                  S 5       rS r\	" S\
S9\	" S5      \	" S5      S	 5       5       5       r\	" S\
S9\	" S5      \	" S5      S
 5       5       5       r\	" S\
S9\	" S5      \	" S5      S 5       5       5       r\	" S\
S9\	" S5      \	" S5      S 5       5       5       r\	" S\
S9\	" S5      S 5       5       r\	" S\
S9\	" S5      S 5       5       rSrg)ÚTestImportJsonFileé   z$Tests for import_json_file function.c                 óL   • SSSSSSSSS	S
SSSSS/SS0S.S	SSSSSSS/0 S./S.$ )zSample JSON data for testing.útest.pdfz2024-01-01T12:00:00é
   é   zTest summaryÚ	anthropiczclaude-3-5-sonnet-20241022r   é   zTable 1zFinancial DatazQuarterly financialsz#| Q1 | Q2 |
|---|---|
| 100 | 200 |ÚQ1ÚQ2ÚkeyÚvalue)Útable_idÚpageÚprocedural_titleÚ	llm_titleÚllm_descriptionÚtable_markdownÚheadersÚmetadatazTable 2z
Sales DatazMonthly salesz#| Jan | Feb |
|---|---|
| 50 | 75 |ÚJanÚFeb)ÚdocumentÚextraction_dateÚtotal_pagesÚtotal_tablesÚexecutive_summaryÚmodel_providerÚmodelÚtables© ©Úselfs    Ú0/home/mande/repo/quber/tests/db/test_importer.pyÚsample_json_dataÚ#TestImportJsonFile.sample_json_data   ss   € ð #Ø4ØØØ!/Ø)Ø1ð !"ØØ(1Ø!1Ø'=Ø&MØ $ d˜|Ø!&¨Ð 0ñ	ð !"ØØ(1Ø!-Ø'6Ø&MØ % u˜~Ø "ñ	ðñ
ð 	
ó    c                 ó€   • [         R                  " [        SS9   [        S5        SSS5        g! , (       d  f       g= f)z7Test that FileNotFoundError is raised for missing file.zJSON file not found©Úmatchz/nonexistent/file.jsonN)ÚpytestÚraisesÚFileNotFoundErrorr
   r,   s    r.   Útest_import_json_file_not_foundÚ2TestImportJsonFile.test_import_json_file_not_found8   ó(   € ä�]Š]Ô,Ð4IÓJÜÐ5Ô6÷ K×JÖJúó   š/¯
=zbuiltins.open)Únew_callableúpathlib.Path.existsúquber.db.importer.get_sessionc                 ó  • SUl         [        R                  " U5      UR                   R                  l         [	        5       nSUR
                  R                   R                  R                   R                  l         [	        US9Ul        [	        SS9Ul	        XQl         [        SUS9   [        S5      u  pgSSS5        WR                  nSo˜U	:H  oª(       dº  [        R                  " SU
4SX‰45      S	[        R                   " 5       ;   d  [        R"                  " U5      (       a  [        R$                  " U5      OS	[        R$                  " U5      [        R$                  " U	5      S
.-  nSSU0-  n['        [        R(                  " U5      5      eS=n=p©UR*                  nSo˜U	:H  oª(       dº  [        R                  " SU
4SX‰45      S	[        R                   " 5       ;   d  [        R"                  " U5      (       a  [        R$                  " U5      OS	[        R$                  " U5      [        R$                  " U	5      S
.-  nSSU0-  n['        [        R(                  " U5      5      eS=n=p©UR,                  nSo˜U	:H  oª(       dº  [        R                  " SU
4SX‰45      S	[        R                   " 5       ;   d  [        R"                  " U5      (       a  [        R$                  " U5      OS	[        R$                  " U5      [        R$                  " U	5      S
.-  nSSU0-  n['        [        R(                  " U5      5      eS=n=p©SnWU:H  oˆ(       d¥  [        R                  " SU4SX}45      S[        R                   " 5       ;   d  [        R"                  " U5      (       a  [        R$                  " U5      OS[        R$                  " U5      S.-  nSSU0-  n['        [        R(                  " U5      5      eS=p�g! , (       d  f       GNJ= f)zTest basic JSON file import.TN©Úreturn_valueú	json.loadú/test/file.jsonr   ©ú==)z0%(py2)s
{%(py2)s = %(py0)s.filename
} == %(py5)sÚdoc©Úpy0Úpy2Úpy5úassert %(py7)sÚpy7r   )z3%(py2)s
{%(py2)s = %(py0)s.total_pages
} == %(py5)sr   )z4%(py2)s
{%(py2)s = %(py0)s.total_tables
} == %(py5)s©z%(py0)s == %(py3)sÚtable_count©rH   Úpy3úassert %(py5)srJ   )rA   ÚjsonÚdumpsÚreadr   ÚqueryÚ	filter_byÚfirstÚ	__enter__Ú__exit__r   r
   ÚfilenameÚ
@pytest_arÚ_call_reprcompareÚ@py_builtinsÚlocalsÚ_should_repr_global_nameÚ	_safereprÚAssertionErrorÚ_format_explanationr%   r&   )r-   Úmock_get_sessionÚmock_existsÚ	mock_filer/   Úmock_sessionrF   rN   Ú@py_assert1Ú@py_assert4Ú@py_assert3Ú@py_format6Ú@py_format8Ú@py_assert2Ú@py_format4s                  r.   Útest_import_json_file_basicÚ.TestImportJsonFile.test_import_json_file_basic=   sN  € ð
 $(ˆÔ Ü37·:²:Ð>NÓ3Oˆ	×Ñ×#Ñ#Ô0ä “{ˆØTXˆ×Ñ×'Ñ'×1Ñ1×>Ñ>×DÑDÔQÜ!*¸Ñ!EˆÔÜ )°tÑ <ˆÔØ(4Ô%ä�;Ð-=Ó>Ü/Ð0AÓBÑˆC÷ ?ð �|‰|Ð)˜zÐ)˜zÑ)×)×)×)Ò)ˆ|×)×)×)×)×)Õ)ˆs×)×)×)Ñ)ˆs×)×)Õ)ˆ|×)×)Ó)˜z×)×)×)×)×)×)Ó)Ø�‰Ð$ "Ð$ "Ñ$×$×$×$Ò$ˆ×$×$×$×$×$Õ$ˆs×$×$×$Ñ$ˆs×$×$Õ$ˆ×$×$Ó$ "×$×$×$×$×$×$Ó$Ø×ÑÐ$ 1Ð$ 1Ñ$×$×$×$Ò$Ð×$×$×$×$×$Õ$ˆs×$×$×$Ñ$ˆs×$×$Õ$Ð×$×$Ó$ 1×$×$×$×$×$×$Ó$ØÐˆ{˜aÑ×××Òˆ{×××××Õˆ{×××Ñˆ{××Õ˜a××××××Ò÷ ?Ö>ús   Â&O7Ï7
Pc                 óŠ  • SUl         [        5       nSUR                  R                   R                  R                   R                  l         [        US9Ul        [        SS9Ul        XQl         [        5       n[        [        S/S-  S9S9UR                  l         [        SUS9   [        SSUSS	9  SSS5        UR                  owR                  nS
o˜U	:H  oª(       dÏ  [        R                  " SU
4SX‰45      S[        R                  " 5       ;   d  [        R                  " U5      (       a  [        R                   " U5      OS[        R                   " U5      [        R                   " U5      [        R                   " U	5      S.-  nSSU0-  n[#        [        R$                  " U5      5      eS=n=n=p©g! , (       d  f       GN	= f)z+Test JSON import with embedding generation.TNr@   gš™™™™™¹?i   )ÚtolistrB   rC   )ÚsessionÚembedding_serviceÚgenerate_embeddingsr   rD   )z\%(py4)s
{%(py4)s = %(py2)s
{%(py2)s = %(py0)s.embed_table_metadata
}.call_count
} == %(py7)sÚmock_embedding_service©rH   rI   Úpy4rL   úassert %(py9)sÚpy9)rA   r   rU   rV   rW   rX   rY   Úembed_table_metadatar   r
   Ú
call_countr[   r\   r]   r^   r_   r`   ra   rb   )r-   rc   rd   re   r/   rf   ru   rg   ri   Ú@py_assert6Ú@py_assert5rk   Ú@py_format10s                r.   Ú%test_import_json_file_with_embeddingsÚ8TestImportJsonFile.test_import_json_file_with_embeddingsS   s?  € ð $(ˆÔ ä “{ˆØTXˆ×Ñ×'Ñ'×1Ñ1×>Ñ>×DÑDÔQÜ!*¸Ñ!EˆÔÜ )°tÑ <ˆÔØ(4Ô%ä!*£ÐÜCLÜ¨3¨%°$©,Ñ7ñD
Ð×3Ñ3Ô@ô �;Ð-=Ó>ÜØ!ØØ"8Ø$(ò	÷ ?ð &×:Ñ:ÐJ×EÑEÐJÈÐJÈÑJ×J×J×JÒJÐE×J×J×J×J×JÕJÐ%×J×J×JÑJÐ%×J×JÕJÐ:×J×JÓJÐE×J×JÓJÈ×J×J×J×J×J×JÖJ÷ ?Ö>ús   Â$F3Æ3
Gc                 óŽ  • SUl         [        SS[        R                  " 5       SS9n[	        5       [	        5       /Ul        [	        5       nXVR                  R                   R                  R                   R                  l         SUR                  R                   R                  R                   R                  l         [	        US9Ul
        [	        SS9Ul        Xal         [        SUS9   [        S	5      u  pxSSS5        WU:H  o™(       dÛ  [        R                  " S
U	4SXu45      S[         R"                  " 5       ;   d  [        R$                  " U5      (       a  [        R&                  " U5      OSS[         R"                  " 5       ;   d  [        R$                  " U5      (       a  [        R&                  " U5      OSS.-  n
SSU
0-  n[)        [        R*                  " U5      5      eSn	SnWU:H  o™(       d¥  [        R                  " S
U	4SXŒ45      S[         R"                  " 5       ;   d  [        R$                  " U5      (       a  [        R&                  " U5      OS[        R&                  " U5      S.-  nSSU0-  n[)        [        R*                  " U5      5      eS=pœUR,                  R/                  5         g! , (       d  f       GNÇ= f)z1Test importing a duplicate file (already exists).Tr   r   r   )ÚidrZ   r$   r&   r@   NrB   rC   rD   )z%(py0)s == %(py2)srF   Úexisting_doc)rH   rI   zassert %(py4)srw   rM   rN   rO   rQ   rJ   )rA   r   r   Únowr   r*   rU   rV   rW   ÚcountrX   rY   r   r
   r[   r\   r]   r^   r_   r`   ra   rb   ÚaddÚassert_not_called)r-   rc   rd   re   r/   rƒ   rf   rF   rN   rg   Ú@py_format3Ú@py_format5rl   rm   rj   s                  r.   Útest_import_json_file_duplicateÚ2TestImportJsonFile.test_import_json_file_duplicater   s¬  € ð
 $(ˆÔ ä 1¨zÌ8Ï<Ê<Ë>ÐhiÑjˆÜ(›{¬I«KÐ8ˆÔä “{ˆØT`×Ñ×'Ñ'×1Ñ1×>Ñ>×DÑDÔQØTUˆ×Ñ×'Ñ'×1Ñ1×>Ñ>×DÑDÔQÜ!*¸Ñ!EˆÔÜ )°tÑ <ˆÔØ(4Ô%ä�;Ð-=Ó>Ü/Ð0AÓBÑˆC÷ ?ð �lÑ"×"×"×"Ò"ˆs×"×"×"×"×"Õ"ˆs×"×"×"Ñ"ˆs×"×"×"×"×"Ö"�l×"×"×"Ñ"�l×"×"×"×"×"×"Ò"ØÐˆ{˜aÑ×××Òˆ{×××××Õˆ{×××Ñˆ{××Õ˜a××××××Ñà×Ñ×*Ñ*Õ,÷ ?Ö>ús   Ã'J5Ê5
Kc                 óô  • SUl         SUS'   [        5       nSUR                  R                   R                  R                   R                  l         [        US9Ul        [        SS9Ul        XQl         [        SUS9   [        S5      u  pgSSS5        WR                  nSo˜U	Loª(       dº  [        R                  " SU
4S	X‰45      S
[        R                  " 5       ;   d  [        R                  " U5      (       a  [        R                  " U5      OS
[        R                  " U5      [        R                  " U	5      S.-  nSSU0-  n[!        [        R"                  " U5      5      eS=n=p©g! , (       d  f       Nå= f)zDTest import with invalid extraction_date falls back to current time.Tzinvalid-dater$   Nr@   rB   rC   )úis not)z;%(py2)s
{%(py2)s = %(py0)s.extraction_date
} is not %(py5)srF   rG   rK   rL   )rA   r   rU   rV   rW   rX   rY   r   r
   r$   r[   r\   r]   r^   r_   r`   ra   rb   )r-   rc   rd   re   r/   rf   rF   Ú_rg   rh   ri   rj   rk   s                r.   Ú"test_import_json_file_invalid_dateÚ5TestImportJsonFile.test_import_json_file_invalid_dateŒ   s   € ð
 $(ˆÔ Ø.<ÐÐ*Ñ+ä “{ˆØTXˆ×Ñ×'Ñ'×1Ñ1×>Ñ>×DÑDÔQÜ!*¸Ñ!EˆÔÜ )°tÑ <ˆÔØ(4Ô%ä�;Ð-=Ó>Ü%Ð&7Ó8‰FˆC÷ ?ð ×"Ñ"Ð.¨$Ð.¨$Ð.×.×.×.Ò.Ð"×.×.×.×.×.Õ.ˆs×.×.×.Ñ.ˆs×.×.Õ.Ð"×.×.Ó.¨$×.×.×.×.×.×.Ô.÷	 ?Õ>ús   Á<E)Å)
E7c                 ó\  • SUl         [        5       nSUR                  R                   R                  R                   R                  l         [        SUS9   [        SUS9  SSS5        UR                  R                  5         UR                  R                  5         g! , (       d  f       NC= f)z$Test import with a provided session.TNrB   r@   rC   ©rr   )
rA   r   rU   rV   rW   r   r
   r†   Úassert_called_onceÚcommit©r-   rd   re   r/   rf   s        r.   Ú+test_import_json_file_with_provided_sessionÚ>TestImportJsonFile.test_import_json_file_with_provided_session    sƒ   € ð $(ˆÔ ä “{ˆØTXˆ×Ñ×'Ñ'×1Ñ1×>Ñ>×DÑDÔQä�;Ð-=Ó>ÜÐ.¸ÒE÷ ?ð 	×Ñ×+Ñ+Ô-Ø×Ñ×.Ñ.Õ0÷ ?Õ>ús   ÁBÂ
B+c                 óÂ  • SUl         [        5       nSUR                  R                   R                  R                   R                  l         [        S5      UR                  l        [        SUS9   [        R                  " [
        SS9   [        SUS9  SSS5        SSS5        UR                  R                  5         g! , (       d  f       N1= f! , (       d  f       N:= f)	z.Test that session errors are handled properly.TNzDatabase errorrB   r@   r3   rC   r’   )rA   r   rU   rV   rW   Ú	Exceptionr”   Úside_effectr   r5   r6   r
   Úrollbackr“   r•   s        r.   Ú#test_import_json_file_session_errorÚ6TestImportJsonFile.test_import_json_file_session_error°   s¤   € ð $(ˆÔ ä “{ˆØTXˆ×Ñ×'Ñ'×1Ñ1×>Ñ>×DÑDÔQÜ*3Ð4DÓ*Eˆ×ÑÔ'ä�;Ð-=Ó>Ü—’œyÐ0@ÓAÜ Ð!2¸LÒI÷ B÷ ?ð 	×Ñ×0Ñ0Õ2÷ BÕAú÷ ?Õ>ús$   Á/CÂ	B?ÂCÂ?
C	Ã	CÃ
Cr+   N)Ú__name__Ú
__module__Ú__qualname__Ú__firstlineno__Ú__doc__r5   Úfixturer/   r8   r   r   rn   r   rŠ   r�   r–   rœ   Ú__static_attributes__r+   r1   r.   r   r      sQ  † Ù.à‡^�^ñ 
ó ð 
òD7ñ
 ˆ?¨Ñ3Ù
Ð Ó!Ù
Ð*Ó+ñ ó ,ó "ó 4ð ñ& ˆ?¨Ñ3Ù
Ð Ó!Ù
Ð*Ó+ñKó ,ó "ó 4ðKñ8 ˆ?¨Ñ3Ù
Ð Ó!Ù
Ð*Ó+ñ-ó ,ó "ó 4ð-ñ. ˆ?¨Ñ3Ù
Ð Ó!Ù
Ð*Ó+ñ/ó ,ó "ó 4ð/ñ" ˆ?¨Ñ3Ù
Ð Ó!ñ1ó "ó 4ð1ñ ˆ?¨Ñ3Ù
Ð Ó!ñ3ó "ó 4ó3r1   r   c                   ó†  • \ rS rSrSrS r\" S5      \" S5      S 5       5       r\" S5      \" S5      \" S5      S 5       5       5       r\" S5      \" S5      \" S5      S	 5       5       5       r	\" S
5      \" S5      \" S5      \" S5      S 5       5       5       5       r
\" S5      \" S5      \" S5      S 5       5       5       rSrg)ÚTestImportDirectoryéÁ   z$Tests for import_directory function.c                 ó€   • [         R                  " [        SS9   [        S5        SSS5        g! , (       d  f       g= f)z<Test that FileNotFoundError is raised for missing directory.zDirectory not foundr3   z/nonexistent/directoryN)r5   r6   r7   r	   r,   s    r.   Útest_import_directory_not_foundÚ3TestImportDirectory.test_import_directory_not_foundÄ   r:   r;   zpathlib.Path.globr=   c                 ó
  • SUl         / Ul         [        S5      u  p4/ oSU:H  of(       d¥  [        R                  " SU4SX545      S[        R
                  " 5       ;   d  [        R                  " U5      (       a  [        R                  " U5      OS[        R                  " U5      S.-  nSSU0-  n[        [        R                  " U5      5      eS	=peS
oTU:H  of(       d¥  [        R                  " SU4SXE45      S[        R
                  " 5       ;   d  [        R                  " U5      (       a  [        R                  " U5      OS[        R                  " U5      S.-  nSSU0-  n[        [        R                  " U5      5      eS	=peg	)z0Test importing directory with no matching files.Tú	/test/dirrD   rM   ÚdocsrO   rQ   rJ   Nr   r&   )
rA   r	   r[   r\   r]   r^   r_   r`   ra   rb   )	r-   rd   Ú	mock_globr­   r&   rl   rg   rm   rj   s	            r.   Útest_import_directory_no_filesÚ2TestImportDirectory.test_import_directory_no_filesÉ   sÖ   € ð $(ˆÔ Ø!#ˆ	Ôä-¨kÓ:ÑˆàÐ�r‰z×××Òˆt×××××Õˆt×××Ñˆt××Õ�r××××××ÑØ Ð ˜qÑ × × × Ò ˆ|× × × × × Õ ˆ|× × × Ñ ˆ|× × Õ ˜q× × × × × × Ò r1   z"quber.db.importer.import_json_filec                 ó  • SUl         [        S5      [        S5      /Ul         [        [        S9n[        [        S9nUS4US4/Ul        [        S5      u  pg[        U5      nSo˜U	:H  oª(       Gd  [        R                  " SU
4S	X‰45      S
[        R                  " 5       ;   d  [        R                  " [        5      (       a  [        R                  " [        5      OS
S[        R                  " 5       ;   d  [        R                  " U5      (       a  [        R                  " U5      OS[        R                  " U5      [        R                  " U	5      S.-  nSSU0-  n[        [        R                  " U5      5      eS=n=p©So‡U:H  oÝ(       d¥  [        R                  " SU4SXx45      S[        R                  " 5       ;   d  [        R                  " U5      (       a  [        R                  " U5      OS[        R                  " U5      S.-  nSSU0-  n[        [        R                  " U5      5      eS=pØUR                  nSo­U
:H  nU(       dº  [        R                  " SU4SXÚ45      S[        R                  " 5       ;   d  [        R                  " U5      (       a  [        R                  " U5      OS[        R                  " U5      [        R                  " U
5      S.-  nSSU0-  n[        [        R                  " U5      5      eS=n=nn
g)z-Test importing multiple files from directory.Tú/test/file1.jsonú/test/file2.json©Úspecr   é   r¬   rD   ©z0%(py3)s
{%(py3)s = %(py0)s(%(py1)s)
} == %(py6)sÚlenr­   ©rH   Úpy1rP   Úpy6úassert %(py8)sÚpy8Né   rM   r&   rO   rQ   rJ   )z2%(py2)s
{%(py2)s = %(py0)s.call_count
} == %(py5)sÚmock_import_filerG   rK   rL   )rA   r   r   r   rš   r	   r¸   r[   r\   r]   r^   r_   r`   ra   rb   r{   )r-   rd   r®   r¿   Ú	mock_doc1Ú	mock_doc2r­   r&   rl   r}   rh   Ú@py_format7Ú@py_format9rg   rm   rj   ri   rk   s                     r.   Ú$test_import_directory_multiple_filesÚ8TestImportDirectory.test_import_directory_multiple_filesÕ   s¾  € ð
 $(ˆÔ Ü"&Ð'9Ó":¼DÐASÓ<TÐ!Uˆ	Ôô ¤8Ñ,ˆ	Ü¤8Ñ,ˆ	Ø)2°A¨¸ÀA¸Ð'GÐÔ$ä-¨kÓ:Ñˆä�4‹yÐ˜AÐ˜A‰~×××Óˆy×××××ÕŒs×××ÑŒs×××××Ö�4×××Ñ�4××Õˆy××Ó˜A××××××ÓØ Ð ˜qÑ × × × Ò ˆ|× × × × × Õ ˆ|× × × Ñ ˆ|× × Õ ˜q× × × × × × Ñ Ø×*Ñ*Ð/¨aÐ/¨aÑ/×/×/×/Ó/Ð*×/×/×/×/×/Õ/Ð×/×/×/Ñ/Ð×/×/Õ/Ð*×/×/Ó/¨a×/×/×/×/×/×/Õ/r1   c                 ó”  • SUl         [        S5      [        S5      [        S5      /Ul         [        [        S9n[        [        S9nUS4[	        S5      US4/Ul        [        S	5      u  pg[        U5      nSo˜U	:H  oª(       Gd  [        R                  " S
U
4SX‰45      S[        R                  " 5       ;   d  [        R                  " [        5      (       a  [        R                  " [        5      OSS[        R                  " 5       ;   d  [        R                  " U5      (       a  [        R                  " U5      OS[        R                  " U5      [        R                  " U	5      S.-  nSSU0-  n[        [        R                  " U5      5      eS=n=p©So‡U:H  oÝ(       d¥  [        R                  " S
U4SXx45      S[        R                  " 5       ;   d  [        R                  " U5      (       a  [        R                  " U5      OS[        R                  " U5      S.-  nSSU0-  n[        [        R                  " U5      5      eS=pØg)z?Test that directory import continues on individual file errors.Tr²   r³   z/test/file3.jsonr´   r   zImport errorr¶   r¬   rD   r·   r¸   r­   r¹   r¼   r½   Nr¾   rM   r&   rO   rQ   rJ   )rA   r   r   r   r™   rš   r	   r¸   r[   r\   r]   r^   r_   r`   ra   rb   )r-   rd   r®   r¿   rÀ   Ú	mock_doc3r­   r&   rl   r}   rh   rÂ   rÃ   rg   rm   rj   s                   r.   Ú!test_import_directory_with_errorsÚ5TestImportDirectory.test_import_directory_with_errorsè   sd  € ð
 $(ˆÔ äÐ#Ó$ÜÐ#Ó$ÜÐ#Ó$ð"
ˆ	Ôô ¤8Ñ,ˆ	Ü¤8Ñ,ˆ	à˜ˆNÜ�nÓ%Ø˜ˆNð(
ÐÔ$ô .¨kÓ:Ñˆô �4‹yÐ˜AÐ˜A‰~×××Óˆy×××××ÕŒs×××ÑŒs×××××Ö�4×××Ñ�4××Õˆy××Ó˜A××××××ÓØ Ð ˜qÑ × × × Ò ˆ|× × × × × Õ ˆ|× × × Ñ ˆ|× × Õ ˜q× × × × × × Ò r1   z'quber.db.importer.get_embedding_servicec                 óH  • SUl         [        S5      /Ul         [        5       nXTl         [        [        S9nUS4Ul         [	        SSS9  UR                  5         UR                  S   nUS   oˆU:H  o™(       d¥  [        R                  " S	U	4S
X…45      [        R                  " U5      S[        R                  " 5       ;   d  [        R                  " U5      (       a  [        R                  " U5      OSS.-  n
SSU
0-  n[        [        R                  " U5      5      eS=p‰US   nSoÈUL o™(       do  [        R                  " SU	4SXŒ45      [        R                  " U5      [        R                  " U5      S.-  nSSU0-  n[        [        R                  " U5      5      eS=n=pœg)z0Test directory import with embedding generation.Tr²   r´   r   r¬   )rt   r   rs   rD   )z%(py1)s == %(py3)sru   )rº   rP   rQ   rJ   Nrt   )Úis)z%(py1)s is %(py4)s©rº   rw   úassert %(py6)sr»   )rA   r   r   r   r	   r“   Ú	call_argsr[   r\   r`   r]   r^   r_   ra   rb   )r-   rd   r®   r¿   Úmock_get_embeddingru   Úmock_docÚcall_kwargsÚ@py_assert0rl   rm   rj   ri   r‰   rÂ   s                  r.   Ú%test_import_directory_with_embeddingsÚ9TestImportDirectory.test_import_directory_with_embeddings  s#  € ð $(ˆÔ Ü"&Ð'9Ó":Ð!;ˆ	Ôä!*£ÐØ*@Ô'ä¤(Ñ+ˆØ)1°1¨ÐÔ%ä˜¸$Ò?ð 	×-Ñ-Ô/ð '×0Ñ0°Ñ3ˆØÐ.Ñ/ÐIÐ3IÑI×I×I×IÒIÐ/×I×IÔIÐ/×I×I×I×I×IÔIÐ3I×I×I×IÑIÐ3I×I×I×I×I×I×IÓIØÐ0Ñ1Ð9°TÐ9°TÐ9×9×9×9Ò9Ð1×9×9Ô9Ð1×9×9Ó9°T×9×9×9×9×9×9Ô9r1   c                 óV   • SUl         / Ul         [        SSS9  UR                  S5        g)z/Test directory import with custom glob pattern.Tr¬   z*.json)ÚpatternN)rA   r	   Úassert_called_once_with)r-   rd   r®   r¿   s       r.   Ú$test_import_directory_custom_patternÚ8TestImportDirectory.test_import_directory_custom_pattern  s.   € ð
 $(ˆÔ Ø!#ˆ	Ôä˜¨hÒ7ð 	×)Ñ)¨(Õ3r1   r+   N)rž   rŸ   r    r¡   r¢   r©   r   r¯   rÄ   rÈ   rÓ   rØ   r¤   r+   r1   r.   r¦   r¦   Á   s  † Ù.ò7ñ
 ÐÓÙ
Ð Ó!ñ!ó "ó  ð!ñ Ð/Ó0Ù
ÐÓÙ
Ð Ó!ñ0ó "ó  ó 1ð0ñ  Ð/Ó0Ù
ÐÓÙ
Ð Ó!ñ!ó "ó  ó 1ð!ñ0 Ð4Ó5Ù
Ð/Ó0Ù
ÐÓÙ
Ð Ó!ñ:ó "ó  ó 1ó 6ð:ñ. Ð/Ó0Ù
ÐÓÙ
Ð Ó!ñ4ó "ó  ó 1ó4r1   r¦   c                   óV   • \ rS rSrSr\" S5      S 5       r\" S5      S 5       rS rSr	g)	ÚTestGetImportStatsi,  z$Tests for get_import_stats function.r>   c                 ó,  • [        5       nSUR                  R                  R                  l        / UR                  R                  R                  l        [        US9Ul        [        SS9Ul        X!l        [        5       nUS   nSoTU:H  of(       do  [        R                  " SU4SXE45      [        R                  " U5      [        R                  " U5      S.-  nSS	U0-  n[        [        R                  " U5      5      eS=n=peUS
   nSoTU:H  of(       do  [        R                  " SU4SXE45      [        R                  " U5      [        R                  " U5      S.-  nSS	U0-  n[        [        R                  " U5      5      eS=n=peUS   nSoTU:H  of(       do  [        R                  " SU4SXE45      [        R                  " U5      [        R                  " U5      S.-  nSS	U0-  n[        [        R                  " U5      5      eS=n=peUS   n/ oTU:H  of(       do  [        R                  " SU4SXE45      [        R                  " U5      [        R                  " U5      S.-  nSS	U0-  n[        [        R                  " U5      5      eS=n=peg)z'Test getting stats from empty database.r   r@   NÚtotal_documentsrD   ©z%(py1)s == %(py4)srÌ   rÍ   r»   r&   Úavg_tables_per_docÚ	documents)r   rU   rA   r…   ÚallrX   rY   r   r[   r\   r`   ra   rb   )	r-   rc   rf   ÚstatsrÒ   ri   rl   r‰   rÂ   s	            r.   Útest_get_import_stats_empty_dbÚ1TestGetImportStats.test_get_import_stats_empty_db/  s©  € ô !“{ˆØ=>ˆ×Ñ×'Ñ'×-Ñ-Ô:Ø;=ˆ×Ñ×'Ñ'×+Ñ+Ô8Ü!*¸Ñ!EˆÔÜ )°tÑ <ˆÔØ(4Ô%ä Ó"ˆàÐ&Ñ'Ð,¨1Ð,¨1Ñ,×,×,×,Ò,Ð'×,×,Ô,Ð'×,×,Ó,¨1×,×,×,×,×,×,Ó,Ø�^Ñ$Ð)¨Ð)¨Ñ)×)×)×)Ò)Ð$×)×)Ô)Ð$×)×)Ó)¨×)×)×)×)×)×)Ó)ØÐ)Ñ*Ð/¨aÐ/¨aÑ/×/×/×/Ò/Ð*×/×/Ô/Ð*×/×/Ó/¨a×/×/×/×/×/×/Ó/Ø�[Ñ!Ð' RÐ' RÑ'×'×'×'Ò'Ð!×'×'Ô'Ð!×'×'Ó' R×'×'×'×'×'×'Ô'r1   c                 óØ	  ^^• [        [        S9mSTl        STl        [        5       [        5       /Tl        [        [        S9mSTl        STl        [        5       [        5       [        5       /Tl        [        5       nUU4S jnX2R
                  l        [        US9Ul        [        SS9Ul        X!l	        [        5       nUS	   nS
oeU:H  ow(       do  [        R                  " SU4SXV45      [        R                  " U5      [        R                  " U5      S.-  nSSU0-  n	[        [        R                  " U	5      5      eS=n=pvUS   nSoeU:H  ow(       do  [        R                  " SU4SXV45      [        R                  " U5      [        R                  " U5      S.-  nSSU0-  n	[        [        R                  " U	5      5      eS=n=pvUS   nSoeU:H  ow(       do  [        R                  " SU4SXV45      [        R                  " U5      [        R                  " U5      S.-  nSSU0-  n	[        [        R                  " U	5      5      eS=n=pvUS   n
[!        U
5      nS
o¶U:H  oÌ(       d×  [        R                  " SU4SXk45      S["        R$                  " 5       ;   d  [        R&                  " [         5      (       a  [        R                  " [         5      OS[        R                  " U
5      [        R                  " U5      [        R                  " U5      S.-  nSSU0-  n[        [        R                  " U5      5      eS=n
=n=pËUS   S   S   nSoeU:H  ow(       do  [        R                  " SU4SXV45      [        R                  " U5      [        R                  " U5      S.-  nSSU0-  n	[        [        R                  " U	5      5      eS=n=pvUS   S   S   nS
oeU:H  ow(       do  [        R                  " SU4SXV45      [        R                  " U5      [        R                  " U5      S.-  nSSU0-  n	[        [        R                  " U	5      5      eS=n=pvUS   S   S   nSoeU:H  ow(       do  [        R                  " SU4SXV45      [        R                  " U5      [        R                  " U5      S.-  nSSU0-  n	[        [        R                  " U	5      5      eS=n=pvg)z)Test getting stats with data in database.r´   zdoc1.pdfr   zdoc2.pdfé   c                 ó²   >• U [         :X  a0  [        5       nSUR                  l        TT/UR                  l        U$ [        5       nSUR                  l        U$ )Nr   r¾   ©r   r   r…   rA   rá   )r)   Ú
mock_queryrÀ   rÁ   s     €€r.   Úquery_side_effectÚMTestGetImportStats.test_get_import_stats_with_data.<locals>.query_side_effectP  sR   ø€ ØœÓ Ü&›[�
Ø01�
× Ñ Ô-Ø/8¸)Ð.D�
—‘Ô+Ø!Ð!ä&›[�
Ø01�
× Ñ Ô-Ø!Ð!r1   r@   NrÝ   r   rD   rÞ   rÌ   rÍ   r»   r&   r¾   rß   g      @rà   )z0%(py4)s
{%(py4)s = %(py0)s(%(py2)s)
} == %(py7)sr¸   rv   rx   ry   r   rZ   r*   r   r¶   )r   r   rZ   r%   r*   rU   rš   rX   rY   rA   r   r[   r\   r`   ra   rb   r¸   r]   r^   r_   )r-   rc   rf   rê   râ   rÒ   ri   rl   r‰   rÂ   rg   r|   r}   rk   r~   rÀ   rÁ   s                  @@r.   Útest_get_import_stats_with_dataÚ2TestGetImportStats.test_get_import_stats_with_data@  s*  ù€ ô ¤8Ñ,ˆ	Ø'ˆ	ÔØ "ˆ	ÔÜ%›K¬«Ð5ˆ	Ôä¤8Ñ,ˆ	Ø'ˆ	ÔØ "ˆ	ÔÜ%›K¬«´i³kÐBˆ	Ôä “{ˆö		"ð *;×ÑÔ&Ü!*¸Ñ!EˆÔÜ )°tÑ <ˆÔØ(4Ô%ä Ó"ˆàÐ&Ñ'Ð,¨1Ð,¨1Ñ,×,×,×,Ò,Ð'×,×,Ô,Ð'×,×,Ó,¨1×,×,×,×,×,×,Ó,Ø�^Ñ$Ð)¨Ð)¨Ñ)×)×)×)Ò)Ð$×)×)Ô)Ð$×)×)Ó)¨×)×)×)×)×)×)Ó)ØÐ)Ñ*Ð1¨cÐ1¨cÑ1×1×1×1Ò1Ð*×1×1Ô1Ð*×1×1Ó1¨c×1×1×1×1×1×1Ó1Ø˜Ñ%Ð+ŒsÐ%Ó&Ð+¨!Ð+¨!Ñ+×+×+×+Ò+Ð&×+×+×+×+×+Õ+Œs×+×+×+Ñ+Œs×+×+Õ+Ð%×+×+Ó+Ð&×+×+Ó+¨!×+×+×+×+×+×+Õ+Ø�[Ñ! !Ñ$ ZÑ0Ð>°JÐ>°JÑ>×>×>×>Ò>Ð0×>×>Ô>Ð0×>×>Ó>°J×>×>×>×>×>×>Ó>Ø�[Ñ! !Ñ$ XÑ.Ð3°!Ð3°!Ñ3×3×3×3Ò3Ð.×3×3Ô3Ð.×3×3Ó3°!×3×3×3×3×3×3Ó3Ø�[Ñ! !Ñ$ XÑ.Ð3°!Ð3°!Ñ3×3×3×3Ò3Ð.×3×3Ô3Ð.×3×3Ó3°!×3×3×3×3×3×3Ô3r1   c                 óÂ  ^	• [        [        S9m	ST	l        ST	l        [        5       /T	l        [        5       nU	4S jnX!R
                  l        [        US9nUS   nSoTU:H  of(       do  [        R                  " SU4S	XE45      [        R                  " U5      [        R                  " U5      S
.-  nSSU0-  n[        [        R                  " U5      5      eS=n=peUS   nSoTU:H  of(       do  [        R                  " SU4S	XE45      [        R                  " U5      [        R                  " U5      S
.-  nSSU0-  n[        [        R                  " U5      5      eS=n=peg)z+Test getting stats with a provided session.r´   r   r¾   c                 ó°   >• U [         :X  a/  [        5       nSUR                  l        T/UR                  l        U$ [        5       nSUR                  l        U$ )Nr   rè   )r)   ré   rÐ   s     €r.   rê   ÚYTestGetImportStats.test_get_import_stats_with_provided_session.<locals>.query_side_effects  sO   ø€ ØœÓ Ü&›[�
Ø01�
× Ñ Ô-Ø/7¨j�
—‘Ô+Ø!Ð!ä&›[�
Ø01�
× Ñ Ô-Ø!Ð!r1   r’   rÝ   r   rD   rÞ   rÌ   rÍ   r»   Nr&   )r   r   rZ   r%   r*   rU   rš   r   r[   r\   r`   ra   rb   )
r-   rf   rê   râ   rÒ   ri   rl   r‰   rÂ   rÐ   s
            @r.   Ú+test_get_import_stats_with_provided_sessionÚ>TestGetImportStats.test_get_import_stats_with_provided_sessionj  sí   ø€ ä¤(Ñ+ˆØ&ˆÔØ ˆÔÜ$›;˜-ˆŒä “{ˆõ		"ð *;×ÑÔ&ä ¨Ñ6ˆàÐ&Ñ'Ð,¨1Ð,¨1Ñ,×,×,×,Ò,Ð'×,×,Ô,Ð'×,×,Ó,¨1×,×,×,×,×,×,Ó,Ø�^Ñ$Ð)¨Ð)¨Ñ)×)×)×)Ò)Ð$×)×)Ô)Ð$×)×)Ó)¨×)×)×)×)×)×)Ô)r1   r+   N)
rž   rŸ   r    r¡   r¢   r   rã   rì   rñ   r¤   r+   r1   r.   rÛ   rÛ   ,  s=   † Ù.á
Ð*Ó+ñ(ó ,ð(ñ  Ð*Ó+ñ'4ó ,ð'4õR*r1   rÛ   )r¢   Úbuiltinsr]   Ú_pytest.assertion.rewriteÚ	assertionÚrewriter[   rR   r   Úpathlibr   Úunittest.mockr   r   r   r5   Úquber.db.importerr   r	   r
   Úquber.db.modelsr   r   r¦   rÛ   r+   r1   r.   Ú<module>rû      sS   ðÙ *ç  „ ƒÝ Ý ß 5Ñ 5ã ÷ñ õ
 %÷l3ñ l3÷^h4ñ h4÷VW*ò W*r1   