test: add zone test + record

genai-impact · Sep 5, 2024 · 40eea0d · 40eea0d
1 parent f3117ee
commit 40eea0d
Show file tree

Hide file tree

Showing 2 changed files with 205 additions and 4 deletions.
diff --git a/tests/cassettes/test_tracer/test_init_with_different_mixes.yaml b/tests/cassettes/test_tracer/test_init_with_different_mixes.yaml
@@ -0,0 +1,170 @@
+interactions:
+- request:
+    body: '{"max_tokens": 1024, "messages": [{"role": "user", "content": "Hello, Claude"}],
+      "model": "claude-3-opus-20240229"}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      anthropic-version:
+      - '2023-06-01'
+      connection:
+      - keep-alive
+      content-length:
+      - '115'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+      user-agent:
+      - Anthropic/Python 0.29.2
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - 'false'
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 0.29.2
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.11.8
+    method: POST
+    uri: https://api.anthropic.com/v1/messages
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAAA1SPQUvDQBCF/8r6Ll62kkZBupeCoLYX8VI9iIQlmYbQzUyancWmIf9dUvXg6cE3
+        773hjWgqOLSxLrLlTpv69W3/9HJ+3z0c9veP5+Nzf4KFDh3NLorR1wSLXsIMfIxNVM8Ki1YqCnAo
+        g08VLW4X0qW4yLP8LsvzFSxKYSVWuI/xr1DpNEcv4rChEOTKbPU6Gm5KMiqmJVIzSLoxG/kypWez
+        NT9fZ2pUKj+sMX1aRJWu6MlHYTgQV4WmnvF7iHRMxCXBcQrBIl2GuBENd0kLlQNxhFtmFpL0H1pN
+        0zcAAAD//wMAXSVRpSYBAAA=
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 8be6aaa4ff78b9b4-BRU
+      Connection:
+      - keep-alive
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 05 Sep 2024 13:46:39 GMT
+      Server:
+      - cloudflare
+      Transfer-Encoding:
+      - chunked
+      X-Robots-Tag:
+      - none
+      anthropic-ratelimit-requests-limit:
+      - '5'
+      anthropic-ratelimit-requests-remaining:
+      - '4'
+      anthropic-ratelimit-requests-reset:
+      - '2024-09-05T13:46:43Z'
+      anthropic-ratelimit-tokens-limit:
+      - '10000'
+      anthropic-ratelimit-tokens-remaining:
+      - '10000'
+      anthropic-ratelimit-tokens-reset:
+      - '2024-09-05T13:46:39Z'
+      request-id:
+      - req_01Bvsh4NeRG6xmqGh3scgoqS
+      via:
+      - 1.1 google
+      x-cloud-trace-context:
+      - 13abff072c4925a394badb55996aee13
+    status:
+      code: 200
+      message: OK
+- request:
+    body: '{"max_tokens": 1024, "messages": [{"role": "user", "content": "Hello, Claude"}],
+      "model": "claude-3-opus-20240229"}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      anthropic-version:
+      - '2023-06-01'
+      connection:
+      - keep-alive
+      content-length:
+      - '115'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+      user-agent:
+      - Anthropic/Python 0.29.2
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - 'false'
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 0.29.2
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.11.8
+    method: POST
+    uri: https://api.anthropic.com/v1/messages
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAAA1SPMUvEQBSE/8o6jc1GklxAbhubE3JibyES1uRxBDf7kry3aAj575JTC6uBb4YZ
+        ZkXfwWGQS5MXRXU4VfX7JMX900Rcvjw/nrzCQpeR9hSJ+AvBYuawAy/Si/q4ZwbuKMChDT51lB0y
+        HpNkZV5WeVkeYdFyVIoK97r+FSp9Xet3cagpBL4xZ70VE/uWjLIZiNQsnO5MzZ+m9dGczc/qTo1y
+        55cHbG8Wojw2M3nhCAeKXaNpjvg1hKZEsSW4mEKwSNcjbkUfx6SN8gdFgStyC076Dx237RsAAP//
+        AwCCH77tJgEAAA==
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 8be6aaad9927b9b4-BRU
+      Connection:
+      - keep-alive
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 05 Sep 2024 13:46:41 GMT
+      Server:
+      - cloudflare
+      Transfer-Encoding:
+      - chunked
+      X-Robots-Tag:
+      - none
+      anthropic-ratelimit-requests-limit:
+      - '5'
+      anthropic-ratelimit-requests-remaining:
+      - '3'
+      anthropic-ratelimit-requests-reset:
+      - '2024-09-05T13:46:43Z'
+      anthropic-ratelimit-tokens-limit:
+      - '10000'
+      anthropic-ratelimit-tokens-remaining:
+      - '10000'
+      anthropic-ratelimit-tokens-reset:
+      - '2024-09-05T13:46:40Z'
+      request-id:
+      - req_01NrSCb7WmP9DgHRPswqwaRk
+      via:
+      - 1.1 google
+      x-cloud-trace-context:
+      - 277d3fc773e2ef9bcd853b88292ce65e
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/test_tracer.py b/tests/test_tracer.py
@@ -2,6 +2,9 @@
 from ecologits import EcoLogits
 from ecologits.exceptions import EcoLogitsError
 
+from openai import OpenAI
+from anthropic import Anthropic
+
 
 @pytest.mark.skip(reason="Double init does not raise an error anymore, but we should test that it works correctly.")
 def test_double_init(tracer_init):
@@ -11,10 +14,7 @@ def test_double_init(tracer_init):
 
 @pytest.mark.skip(reason="Must implement un-instrument behavior first.")
 def test_init_openai_only():
-    from openai import OpenAI
-    from anthropic import Anthropic
-
-    EcoLogits.init("openai")
+    EcoLogits.init(providers=["openai"])
 
     openai_client = OpenAI()
     openai_response = openai_client.chat.completions.create(
@@ -32,3 +32,34 @@ def test_init_openai_only():
     )
     assert len(anthropic_response.content) > 0
     assert not hasattr(anthropic_response, "impacts")
+
+@pytest.mark.vcr
+def test_init_with_different_mixes():
+    client = Anthropic()
+
+    EcoLogits.init(electricity_mix_zone="FRA")
+    message_under_french_mix = client.messages.create(
+        max_tokens=1024,
+        messages=[
+            {
+                "role": "user",
+                "content": "Hello, Claude",
+            }
+        ],
+        model="claude-3-opus-20240229",
+    )
+
+    EcoLogits.init(electricity_mix_zone="BEL")
+    message_under_belgian_mix = client.messages.create(
+        max_tokens=1024,
+        messages=[
+            {
+                "role": "user",
+                "content": "Hello, Claude",
+            }
+        ],
+        model="claude-3-opus-20240229",
+    )
+
+    assert message_under_belgian_mix.impacts.energy.value.min != message_under_french_mix.impacts.energy.value.min
+