From 9471deaaa62a9c19c26f304a2bfe0a64bdf20ec4 Mon Sep 17 00:00:00 2001
From: Stainless Bot <dev@stainlessapi.com>
Date: Wed, 15 Nov 2023 16:53:32 +0000
Subject: [PATCH] feat(client): support reading the base url from an env
 variable

---
 README.md            |  1 +
 src/orb/_client.py   |  4 ++++
 tests/test_client.py | 12 ++++++++++++
 tests/utils.py       | 17 ++++++++++++++++-
 4 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 1f7966df..fa89325c 100644
--- a/README.md
+++ b/README.md
@@ -314,6 +314,7 @@ import httpx
 from orb import Orb
 
 client = Orb(
+    # Or use the `ORB_BASE_URL` env var
     base_url="http://my.test.server.example.com:8083",
     http_client=httpx.Client(
         proxies="http://my.test.proxy.example.com",
diff --git a/src/orb/_client.py b/src/orb/_client.py
index 75e6426e..06899550 100644
--- a/src/orb/_client.py
+++ b/src/orb/_client.py
@@ -90,6 +90,8 @@ def __init__(
             )
         self.api_key = api_key
 
+        if base_url is None:
+            base_url = os.environ.get("ORB_BASE_URL")
         if base_url is None:
             base_url = f"https://api.withorb.com/v1"
 
@@ -326,6 +328,8 @@ def __init__(
             )
         self.api_key = api_key
 
+        if base_url is None:
+            base_url = os.environ.get("ORB_BASE_URL")
         if base_url is None:
             base_url = f"https://api.withorb.com/v1"
 
diff --git a/tests/test_client.py b/tests/test_client.py
index 0c73d2b1..1d1afe52 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -25,6 +25,8 @@
     make_request_options,
 )
 
+from .utils import update_env
+
 base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
 api_key = "My API Key"
 
@@ -421,6 +423,11 @@ def test_idempotency_header_options(self, respx_mock: MockRouter) -> None:
         )
         assert response.request.headers.get("Idempotency-Key") == "custom-key"
 
+    def test_base_url_env(self) -> None:
+        with update_env(ORB_BASE_URL="http://localhost:5000/from/env"):
+            client = Orb(api_key=api_key, _strict_response_validation=True)
+            assert client.base_url == "http://localhost:5000/from/env/"
+
     @pytest.mark.parametrize(
         "client",
         [
@@ -975,6 +982,11 @@ async def test_idempotency_header_options(self, respx_mock: MockRouter) -> None:
         )
         assert response.request.headers.get("Idempotency-Key") == "custom-key"
 
+    def test_base_url_env(self) -> None:
+        with update_env(ORB_BASE_URL="http://localhost:5000/from/env"):
+            client = AsyncOrb(api_key=api_key, _strict_response_validation=True)
+            assert client.base_url == "http://localhost:5000/from/env/"
+
     @pytest.mark.parametrize(
         "client",
         [
diff --git a/tests/utils.py b/tests/utils.py
index 766c25c5..ade38898 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import os
 import traceback
-from typing import Any, TypeVar, cast
+import contextlib
+from typing import Any, TypeVar, Iterator, cast
 from datetime import date, datetime
 from typing_extensions import Literal, get_args, get_origin, assert_type
 
@@ -103,3 +105,16 @@ def _assert_list_type(type_: type[object], value: object) -> None:
     inner_type = get_args(type_)[0]
     for entry in value:
         assert_type(inner_type, entry)  # type: ignore
+
+
+@contextlib.contextmanager
+def update_env(**new_env: str) -> Iterator[None]:
+    old = os.environ.copy()
+
+    try:
+        os.environ.update(new_env)
+
+        yield None
+    finally:
+        os.environ.clear()
+        os.environ.update(old)