doc-intel-api/fix_test2.py at master · DavidEscotoDev/doc-intel-api · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
f=open('tests/test_middleware.py')
c=f.read()
f.close()

old = """    @pytest.mark.asyncio
    async def test_different_keys_have_separate_limits(self, client: AsyncClient, auth_headers: dict, test_api_key, db_session, test_document):
        \"\"\"Test that different API keys have separate rate limits.\"\"\"
        from app.models.api_key import APIKey
        from passlib.context import CryptContext

        pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
        other_key = APIKey(
            key_hash=pwd_context.hash("di_other_key_1234567890123456"),
            name="Other Key",
            rate_limit=100,
            is_active=True,
        )
        db_session.add(other_key)
        await db_session.commit()
        await db_session.refresh(other_key)

        other_headers = {"Authorization": f"Bearer di_other_key_1234567890123456"}

        # Use up limit on first key
        for _ in range(99):
            response = await client.get(
                f"/api/v1/process/status/{test_document.id}",
                headers=auth_headers,
            )
            assert response.status_code == status.HTTP_200_OK

        # Other key should still have full limit
        response = await client.get(
            f"/api/v1/process/status/{test_document.id}",
            headers=other_headers,
        )
        assert response.status_code == status.HTTP_200_OK
        remaining = int(response.headers["x-rate-limit-remaining"])
        assert remaining >= 90  # Should have most of its limit left"""

new = """    @pytest.mark.asyncio
    async def test_different_keys_have_separate_limits(self, client: AsyncClient, auth_headers: dict, test_api_key, db_session, test_document):
        \"\"\"Test that different API keys have separate rate limits.\"\"\"
        from app.models.api_key import APIKey
        from app.models.document import Document, DocumentStatus
        from passlib.context import CryptContext

        pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
        other_key = APIKey(
            key_hash=pwd_context.hash("di_other_key_1234567890123456"),
            name="Other Key",
            rate_limit=100,
            is_active=True,
        )
        db_session.add(other_key)
        await db_session.commit()
        await db_session.refresh(other_key)

        # Create a document for the other key
        other_doc = Document(
            filename="other.pdf",
            file_path="/tmp/other.pdf",
            mime_type="application/pdf",
            file_size=1024,
            status=DocumentStatus.UPLOADED,
            api_key_id=other_key.id,
        )
        db_session.add(other_doc)
        await db_session.commit()
        await db_session.refresh(other_doc)

        other_headers = {"Authorization": f"Bearer di_other_key_1234567890123456"}

        # Use up limit on first key
        for _ in range(99):
            response = await client.get(
                f"/api/v1/process/status/{test_document.id}",
                headers=auth_headers,
            )
            assert response.status_code == status.HTTP_200_OK

        # Other key should still have full limit
        response = await client.get(
            f"/api/v1/process/status/{other_doc.id}",
            headers=other_headers,
        )
        assert response.status_code == status.HTTP_200_OK
        remaining = int(response.headers["x-rate-limit-remaining"])
        assert remaining >= 90  # Should have most of its limit left"""

c = c.replace(old, new)
f=open('tests/test_middleware.py','w')
f.write(c)
f.close()
print('Done')