Skip to content

Commit

Permalink
Fixes left out of #1589 (#1593)
Browse files Browse the repository at this point in the history
## Description

Test and transaction updates to get integration tests passing again.

- [ ] I have reviewed the [Guidelines for Contributing](CONTRIBUTING.md)
and the [Code of Conduct](CODE_OF_CONDUCT.md).
  • Loading branch information
richard-rogers authored Dec 3, 2024
1 parent f123399 commit 1d975e7
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 6 deletions.
9 changes: 4 additions & 5 deletions python/tests/api/writer/test_whylabs_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,8 @@ def _get_org() -> str:
@httpretty.activate(allow_net_connect=False, verbose=True)
def test_whylabs_writer_throttle_retry():
ENDPOINT = os.environ["WHYLABS_API_ENDPOINT"]
ORG_ID = _get_org()
MODEL_ID = "XXX"
uri = f"{ENDPOINT}/v0/organizations/{ORG_ID}/log/async/{MODEL_ID}"
uri = f"{ENDPOINT}/v1/log/async"
httpretty.register_uri(httpretty.POST, uri, status=429) # Fake WhyLabs that throttles
why.init(reinit=True, force_local=True)
data = {"col1": 1, "col2": "foo"}
Expand Down Expand Up @@ -271,9 +270,8 @@ def test_put_column_schema_retry():
@httpretty.activate(allow_net_connect=False, verbose=True)
def test_log_async_retry():
ENDPOINT = os.environ["WHYLABS_API_ENDPOINT"]
ORG_ID = _get_org()
MODEL_ID = "xxx"
uri = f"{ENDPOINT}/v0/organizations/{ORG_ID}/log/async/{MODEL_ID}"
uri = f"{ENDPOINT}/v1/log/async"
httpretty.register_uri(httpretty.POST, uri, status=429) # Fake WhyLabs that throttles

writer = WhyLabsWriter(dataset_id=MODEL_ID)
Expand Down Expand Up @@ -651,8 +649,9 @@ def test_transaction_aborted():
transaction_id = writer.start_transaction()
status, id = writer.write(result)
assert status
writer._whylabs_client.abort_transaction(transaction_id)
writer.abort_transaction()
status, id = writer.write(result)
assert transaction_id == writer._whylabs_client._transaction_id
with pytest.raises(TransactionAbortedException) as e:
writer.commit_transaction()
assert str(e) == "Transaction has been aborted"
Expand Down
2 changes: 1 addition & 1 deletion python/whylogs/api/writer/whylabs_transaction_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ def __init__(
transaction_id: Optional[str] = None,
):
super().__init__(org_id, api_key, dataset_id, api_client, ssl_ca_cert, _timeout_seconds, whylabs_client)
transaction_id = transaction_id or self._whylabs_client.get_transaction_id() # type: ignore
transaction_id = transaction_id or self._whylabs_client._transaction_id or self._whylabs_client.get_transaction_id() # type: ignore
self._whylabs_client._transaction_id = transaction_id # type: ignore
self._aborted: bool = False

Expand Down

0 comments on commit 1d975e7

Please sign in to comment.