Repository: abhinavsingh/proxy.py
Pull request #1510 (2024-11-22 04:21:26)

Leakage class

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) Report
Attention: Patch coverage is `92.85714%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 84.08%. Comparing base [(`71d796b`)](https://app.codecov.io/gh/abhinavsingh/proxy.py/commit/71d796baa961cacda7c1e661a98f10e1e474cd85?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) to head [(`c771734`)](https://app.codecov.io/gh/abhinavsingh/proxy.py/commit/c771734ea5257891c4dc737cd2f2af43002ec215?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).

| [Files with missing lines](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | Patch % | Lines |
|---|---|---|
| [proxy/common/leakage.py](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510?src=pr&el=tree&filepath=proxy%2Fcommon%2Fleakage.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29tbW9uL2xlYWthZ2UucHk=) | 82.60% | [4 Missing :warning: ](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           develop    #1510      +/-   ##
===========================================
- Coverage    84.58%   84.08%   -0.51%     
===========================================
  Files          184      186       +2     
  Lines         8364     8420      +56     
  Branches      1276     1279       +3     
===========================================
+ Hits          7075     7080       +5     
- Misses        1083     1127      +44     
- Partials       206      213       +7     
```

| [Flag](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | Coverage Δ | |
|---|---|---|
| [GHA](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `84.06% <92.85%> (-0.39%)` | :arrow_down: |
| [Linux](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `84.06% <92.85%> (-0.03%)` | :arrow_down: |
| [Python 3.10](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `85.18% <92.85%> (+0.05%)` | :arrow_up: |
| [Python 3.11](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `79.64% <92.85%> (-4.71%)` | :arrow_down: |
| [Python 3.12](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `?` | |
| [Python 3.7](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `?` | |
| [Python 3.8](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `85.18% <92.85%> (+0.05%)` | :arrow_up: |
| [Python 3.9](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `?` | |
| [Windows](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `?` | |
| [macOS](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `?` | |
| [pytest](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | `84.06% <92.85%> (-0.39%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/abhinavsingh/proxy.py/pull/1510?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).

----
🚨 Try these New Features:

- [Flaky Tests Detection](https://docs.codecov.com/docs/test-result-ingestion-beta) - Detect and resolve failed and flaky tests
- [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis) - Avoid shipping oversized bundles

Patch:
diff --git a/proxy/common/leakage.py b/proxy/common/leakage.py
new file mode 100644
index 0000000..914b76f
--- /dev/null
+++ b/proxy/common/leakage.py
@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+import time
+
+
+class Leakage:
+    """Leaky Bucket algorithm."""
+
+    def __init__(self, rate: int) -> None:
+        """Initialize the leaky bucket with a specified leak rate in bytes per second."""
+        # Maximum number of tokens the bucket can hold (bytes per second)
+        self.rate = rate
+        self.tokens = rate
+        self.last_check = time.time()
+
+    def _refill(self) -> None:
+        """Refill tokens based on the elapsed time since the last check."""
+        now = time.time()
+        elapsed = now - self.last_check
+        # Add tokens proportional to elapsed time, up to the rate
+        self.tokens += int(elapsed * self.rate)
+        # Cap tokens at the maximum rate to enforce the rate limit
+        self.tokens = min(self.tokens, self.rate)
+        self.last_check = now
+
+    def release(self, tokens: int) -> None:
+        """When you are unable to consume amount units of token, release them into the bucket.
+
+        E.g. say you wanted to read 1024 units, but only 24 units were read, then put
+        back unconsumed 1000 tokens back in the bucket."""
+        if tokens < 0:
+            raise ValueError('Cannot release a negative number of tokens')
+        self.tokens += tokens
+        self.tokens = min(self.tokens, self.rate)
+
+    def consume(self, amount: int) -> int:
+        """Attempt to consume the amount from the bucket.
+
+        Returns the amount allowed to be sent, up to the available tokens (rate).
+        """
+        self._refill()
+        allowed = min(amount, self.tokens)
+        self.tokens -= allowed
+        return allowed


Test patch:
diff --git a/tests/common/test_leakage.py b/tests/common/test_leakage.py
new file mode 100644
index 0000000..564139c
--- /dev/null
+++ b/tests/common/test_leakage.py
@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+import time
+
+import unittest
+
+from proxy.common.leakage import Leakage
+
+
+class TestLeakage(unittest.TestCase):
+
+    def test_initial_consume_no_tokens(self) -> None:
+        # Test consuming with no tokens available initially
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+        self.assertEqual(
+            bucket.consume(150),
+            100,
+        )  # No tokens yet, so expect 0 bytes to be sent
+
+    def test_consume_with_refill(self) -> None:
+        # Test consuming with refill after waiting
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+        time.sleep(1)  # Wait for a second to allow refill
+        self.assertEqual(bucket.consume(50), 50)  # 50 bytes should be available
+
+    def test_consume_above_leak_rate(self) -> None:
+        # Test attempting to consume more than the leak rate after a refill
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+        time.sleep(1)  # Wait for a second to allow refill
+        self.assertEqual(bucket.consume(150), 100)  # Only 100 bytes should be allowed
+
+    def test_repeated_consume_with_partial_refill(self) -> None:
+        # Test repeated consumption with partial refill
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+
+        time.sleep(1)  # Allow tokens to accumulate
+        bucket.consume(80)  # Consume 80 bytes, should leave 20
+        time.sleep(0.5)  # Wait half a second to refill by 50 bytes
+
+        self.assertEqual(bucket.consume(50), 50)  # 50 bytes should be available now
+
+    def test_negative_token_guard(self) -> None:
+        # Ensure tokens do not go negative
+        rate = 100  # bytes per second
+        bucket = Leakage(rate)
+        time.sleep(1)  # Allow tokens to accumulate
+        bucket.consume(150)  # Consume all available tokens
+        self.assertEqual(bucket.consume(10), 0)  # Should return 0 as no tokens are left
+        self.assertEqual(bucket.tokens, 0)  # Tokens should not be negative