diff --git a/.github/workflows/pytests.yml b/.github/workflows/pytests.yml index 4cf5ed85..fcd922ec 100644 --- a/.github/workflows/pytests.yml +++ b/.github/workflows/pytests.yml @@ -22,6 +22,7 @@ jobs: GRPC_PORT: ${{ secrets.GRPC_PORT }} LCD_PORT: ${{ secrets.LCD_PORT }} NETWORK_INSECURE: ${{ secrets.NETWORK_INSECURE }} + WEBSOCKET_ENDPOINT: ${{ secrets.WEBSOCKET_ENDPOINT }} steps: # ---------------------------------------------- # check-out repo and set-up python diff --git a/nibiru/__init__.py b/nibiru/__init__.py index ca758d38..ac9490bb 100644 --- a/nibiru/__init__.py +++ b/nibiru/__init__.py @@ -11,8 +11,9 @@ import nibiru.common # noqa +import nibiru.msg # noqa from nibiru.client import GrpcClient # noqa -from nibiru.common import Coin, Direction, PoolAsset, Side, TxConfig # noqa +from nibiru.common import Coin, Direction, PoolAsset, Side, TxConfig, TxType # noqa from nibiru.network import Network # noqa from nibiru.sdk import Sdk # noqa from nibiru.transaction import Transaction # noqa diff --git a/nibiru/event_specs.py b/nibiru/event_specs.py new file mode 100644 index 00000000..ecb58617 --- /dev/null +++ b/nibiru/event_specs.py @@ -0,0 +1,28 @@ +from dataclasses import dataclass +from enum import Enum + + +class EventType(Enum): + """ + The events enum type shows the type of events available to parse from the nibiruwebsocket object. + """ + + # Perp events + PositionChangedEvent = "nibiru.perp.v1.PositionChangedEvent" + PositionLiquidatedEvent = "nibiru.perp.v1.PositionLiquidatedEvent" + FundingRateChangedEvent = "nibiru.perp.v1.FundingRateChangedEvent" + PositionSettledEvent = "nibiru.perp.v1.PositionSettledEvent" + + # Vpool events + ReserveSnapshotSavedEvent = "nibiru.vpool.v1.ReserveSnapshotSavedEvent" + SwapQuoteForBaseEvent = "nibiru.vpool.v1.SwapQuoteForBaseEvent" + MarkPriceChanged = "nibiru.vpool.v1.MarkPriceChanged" + + # Bank + Transfer = "transfer" + + +@dataclass +class EventCaptured: + event_type: EventType + payload: dict diff --git a/nibiru/network.py b/nibiru/network.py index 2c0ba1b8..a20e7303 100644 --- a/nibiru/network.py +++ b/nibiru/network.py @@ -17,6 +17,7 @@ class Network: chain_id: str fee_denom: str env: str + websocket_endpoint: str = None def __post_init__(self): """ @@ -55,6 +56,7 @@ def devnet(cls) -> "Network": return cls( lcd_endpoint=f'http://{chain_config["HOST"]}:{chain_config["LCD_PORT"]}', grpc_endpoint=f'{chain_config["HOST"]}:{chain_config["GRPC_PORT"]}', + websocket_endpoint=os.getenv("WEBSOCKET_ENDPOINT"), chain_id=chain_config["CHAIN_ID"], fee_denom='unibi', env="devnet", diff --git a/nibiru/utils.py b/nibiru/utils.py index e0947c7d..687a7443 100644 --- a/nibiru/utils.py +++ b/nibiru/utils.py @@ -1,3 +1,5 @@ +import logging +import sys from datetime import datetime from typing import Any, Callable, Union @@ -175,3 +177,68 @@ def toPbTimestamp(dt: datetime): ts = Timestamp() ts.FromDatetime(dt) return ts + + +class ColoredFormatter(logging.Formatter): + + fmt = "%(asctime)s|%(levelname)s|%(funcName)s| %(message)s" + + white = "\x1b[97;20m" + grey = "\x1b[38;20m" + green = "\x1b[32;20m" + cyan = "\x1b[36;20m" + yellow = "\x1b[33;20m" + red = "\x1b[31;20m" + bold_red = "\x1b[31;1m" + reset = "\x1b[0m" + + FORMATS = { + logging.DEBUG: fmt.format(green, reset), + logging.INFO: fmt.format(cyan, reset), + logging.WARNING: fmt.format(yellow, reset), + logging.ERROR: fmt.format(red, reset), + logging.CRITICAL: fmt.format(bold_red, reset), + } + + def format(self, record: logging.LogRecord): + """Formats a record for the logging handler. + + Args: + record (logging.LogRecord): Represents an instance of an event being + logged. + """ + log_format = self.FORMATS.get(record.levelno) + formatter = logging.Formatter(log_format, datefmt="%H:%M:%S") + return formatter.format(record=record) + + +def init_logger(name: str) -> logging.Logger: + """ + Simple logger to use throughout the test suite. + + Examples: + ```python + from nibiru.utils import init_logger + LOGGER = init_logger("test-logger") + LOGGER.info("successfully executed tx staking command") + LOGGER.debug("debugging error message") + ``` + + Log levels include: [debug, info, warning, error, critical] + + Args: + name (str): Name of the logger + + Returns: + logging.Logger: The logger object + """ + logger = logging.getLogger(name) + logger.setLevel(logging.DEBUG) + + # Logs to stdout so we can at least see logs in GHA. + handler = logging.StreamHandler(sys.stdout) + handler.setLevel(logging.DEBUG) + + handler.setFormatter(fmt=ColoredFormatter()) + logger.addHandler(handler) + return logger diff --git a/nibiru/websocket.py b/nibiru/websocket.py new file mode 100644 index 00000000..733c30c6 --- /dev/null +++ b/nibiru/websocket.py @@ -0,0 +1,115 @@ +import json +import threading +import time +from multiprocessing import Queue +from typing import List + +from websocket import WebSocketApp + +from nibiru import Network +from nibiru.event_specs import EventCaptured, EventType +from nibiru.utils import init_logger + +ERROR_TIMEOUT_SLEEP = 3 + + +class NibiruWebsocket: + queue: Queue = None + captured_events_type: List[List[str]] + + def __init__( + self, + network: Network, + captured_events_type: List[EventType] = [], + ): + """ + The nibiru listener provides an interface to easily connect and handle subscription to the events of a nibiru + chain. + """ + + self.websocket_url = network.websocket_endpoint + if self.websocket_url is None: + raise ValueError( + "No websocket endpoint provided. Construct the network object setting up the " + "`websocket_endpoint` endpoint" + ) + + self.captured_events_type: dict = { + captured_event.value: captured_event + for captured_event in captured_events_type + } + self.queue = Queue() + self.logger = init_logger("ws-logger") + + def start(self): + """ + Start the websocket and fill the queue with events. + """ + + self._ws = WebSocketApp( + self.websocket_url, + on_open=self._on_open, + on_message=self._on_message, + on_error=self._on_error, + ) + threading.Thread( + target=self._ws.run_forever, + daemon=True, + name=f"Nibiru websocket @ {self.websocket_url}", + ).start() + + def _on_open(self, _: WebSocketApp): + self.logger.info("WebSocket starting") + self._subscribe() + + def _on_error(self, ws: WebSocketApp, error: Exception): + self.logger.error(f"Closing websocket, error {error}") + self.logger.exception(error) + ws.close() + time.sleep(ERROR_TIMEOUT_SLEEP) + ws.run_forever() + + def _on_message(self, _: WebSocketApp, message: str): + """ + Parse the message and filter through them using the captured event type. + Put these filtered event in the queue. + + Args: + _ (WebSocketApp): No idea what this is + message (str): The message in a utf-8 data received from the server + """ + log = json.loads(message).get("result") + if log is None: + return + + events = log.get("events") + if events is None: + return + + block_height = int(log["data"]["value"]["TxResult"]["height"]) + tx_hash = events["tx.hash"][0] + + events = json.loads(log["data"]["value"]["TxResult"]["result"]["log"])[0] + + for event in events["events"]: + if event["type"] in self.captured_events_type: + event_payload = { + attribute["key"]: attribute["value"].strip('"') + for attribute in event["attributes"] + } + event_payload["block_height"] = block_height + event_payload["tx_hash"] = tx_hash + + self.queue.put(EventCaptured(event["type"], event_payload)) + + def _subscribe(self): + self._ws.send( + json.dumps( + { + "jsonrpc": "2.0", + "method": "subscribe", + "id": 1, + "params": {"query": "tm.event='Tx'"}, + } + ) + ) diff --git a/poetry.lock b/poetry.lock index c815703a..6bc59839 100644 --- a/poetry.lock +++ b/poetry.lock @@ -123,7 +123,7 @@ coincurve = ">=15.0,<18" [[package]] name = "black" -version = "22.6.0" +version = "22.8.0" description = "The uncompromising code formatter." category = "dev" optional = false @@ -500,11 +500,11 @@ pyparsing = ">=2.0.2,<3.0.5 || >3.0.5" [[package]] name = "pathspec" -version = "0.9.0" +version = "0.10.0" description = "Utility library for gitignore style pattern matching of file paths." category = "dev" optional = false -python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,>=2.7" +python-versions = ">=3.7" [[package]] name = "platformdirs" @@ -780,7 +780,7 @@ socks = ["PySocks (>=1.5.6,!=1.5.7,<2.0)"] [[package]] name = "virtualenv" -version = "20.16.3" +version = "20.16.4" description = "Virtual Python Environment builder" category = "main" optional = false @@ -795,6 +795,19 @@ platformdirs = ">=2.4,<3" docs = ["proselint (>=0.13)", "sphinx (>=5.1.1)", "sphinx-argparse (>=0.3.1)", "sphinx-rtd-theme (>=1)", "towncrier (>=21.9)"] testing = ["coverage (>=6.2)", "coverage-enable-subprocess (>=1)", "flaky (>=3.7)", "packaging (>=21.3)", "pytest (>=7.0.1)", "pytest-env (>=0.6.2)", "pytest-freezegun (>=0.4.2)", "pytest-mock (>=3.6.1)", "pytest-randomly (>=3.10.3)", "pytest-timeout (>=2.1)"] +[[package]] +name = "websocket-client" +version = "1.4.0" +description = "WebSocket client for Python with low level API options" +category = "main" +optional = false +python-versions = ">=3.7" + +[package.extras] +test = ["websockets"] +optional = ["wsaccel", "python-socks"] +docs = ["sphinx-rtd-theme (>=0.5)", "Sphinx (>=3.4)"] + [[package]] name = "yarl" version = "1.8.1" @@ -810,7 +823,7 @@ multidict = ">=4.0" [metadata] lock-version = "1.1" python-versions = "^3.9" -content-hash = "60f20cc223f5ba9d39a51bfef0c67ec26adc16e28c44ce8a6395ab76ca287566" +content-hash = "9de97a705b2b8b01d03d1492f167d66cd378da3357f01fa832a8d0a0722a7119" [metadata.files] aiocron = [ @@ -929,29 +942,29 @@ bip32 = [ {file = "bip32-3.1.linux-x86_64.tar.gz", hash = "sha256:9736aa445eb6fa45c752a65d0861ad813c2d7c390b3ee920e3ddc11e3dfa93f7"}, ] black = [ - {file = "black-22.6.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:f586c26118bc6e714ec58c09df0157fe2d9ee195c764f630eb0d8e7ccce72e69"}, - {file = "black-22.6.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b270a168d69edb8b7ed32c193ef10fd27844e5c60852039599f9184460ce0807"}, - {file = "black-22.6.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6797f58943fceb1c461fb572edbe828d811e719c24e03375fd25170ada53825e"}, - {file = "black-22.6.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c85928b9d5f83b23cee7d0efcb310172412fbf7cb9d9ce963bd67fd141781def"}, - {file = "black-22.6.0-cp310-cp310-win_amd64.whl", hash = "sha256:f6fe02afde060bbeef044af7996f335fbe90b039ccf3f5eb8f16df8b20f77666"}, - {file = "black-22.6.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:cfaf3895a9634e882bf9d2363fed5af8888802d670f58b279b0bece00e9a872d"}, - {file = "black-22.6.0-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:94783f636bca89f11eb5d50437e8e17fbc6a929a628d82304c80fa9cd945f256"}, - {file = "black-22.6.0-cp36-cp36m-win_amd64.whl", hash = "sha256:2ea29072e954a4d55a2ff58971b83365eba5d3d357352a07a7a4df0d95f51c78"}, - {file = "black-22.6.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:e439798f819d49ba1c0bd9664427a05aab79bfba777a6db94fd4e56fae0cb849"}, - {file = "black-22.6.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:187d96c5e713f441a5829e77120c269b6514418f4513a390b0499b0987f2ff1c"}, - {file = "black-22.6.0-cp37-cp37m-win_amd64.whl", hash = "sha256:074458dc2f6e0d3dab7928d4417bb6957bb834434516f21514138437accdbe90"}, - {file = "black-22.6.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:a218d7e5856f91d20f04e931b6f16d15356db1c846ee55f01bac297a705ca24f"}, - {file = "black-22.6.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:568ac3c465b1c8b34b61cd7a4e349e93f91abf0f9371eda1cf87194663ab684e"}, - {file = "black-22.6.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:6c1734ab264b8f7929cef8ae5f900b85d579e6cbfde09d7387da8f04771b51c6"}, - {file = "black-22.6.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c9a3ac16efe9ec7d7381ddebcc022119794872abce99475345c5a61aa18c45ad"}, - {file = "black-22.6.0-cp38-cp38-win_amd64.whl", hash = "sha256:b9fd45787ba8aa3f5e0a0a98920c1012c884622c6c920dbe98dbd05bc7c70fbf"}, - {file = "black-22.6.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:7ba9be198ecca5031cd78745780d65a3f75a34b2ff9be5837045dce55db83d1c"}, - {file = "black-22.6.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:a3db5b6409b96d9bd543323b23ef32a1a2b06416d525d27e0f67e74f1446c8f2"}, - {file = "black-22.6.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:560558527e52ce8afba936fcce93a7411ab40c7d5fe8c2463e279e843c0328ee"}, - {file = "black-22.6.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b154e6bbde1e79ea3260c4b40c0b7b3109ffcdf7bc4ebf8859169a6af72cd70b"}, - {file = "black-22.6.0-cp39-cp39-win_amd64.whl", hash = "sha256:4af5bc0e1f96be5ae9bd7aaec219c901a94d6caa2484c21983d043371c733fc4"}, - {file = "black-22.6.0-py3-none-any.whl", hash = "sha256:ac609cf8ef5e7115ddd07d85d988d074ed00e10fbc3445aee393e70164a2219c"}, - {file = "black-22.6.0.tar.gz", hash = "sha256:6c6d39e28aed379aec40da1c65434c77d75e65bb59a1e1c283de545fb4e7c6c9"}, + {file = "black-22.8.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:ce957f1d6b78a8a231b18e0dd2d94a33d2ba738cd88a7fe64f53f659eea49fdd"}, + {file = "black-22.8.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5107ea36b2b61917956d018bd25129baf9ad1125e39324a9b18248d362156a27"}, + {file = "black-22.8.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e8166b7bfe5dcb56d325385bd1d1e0f635f24aae14b3ae437102dedc0c186747"}, + {file = "black-22.8.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd82842bb272297503cbec1a2600b6bfb338dae017186f8f215c8958f8acf869"}, + {file = "black-22.8.0-cp310-cp310-win_amd64.whl", hash = "sha256:d839150f61d09e7217f52917259831fe2b689f5c8e5e32611736351b89bb2a90"}, + {file = "black-22.8.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:a05da0430bd5ced89176db098567973be52ce175a55677436a271102d7eaa3fe"}, + {file = "black-22.8.0-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4a098a69a02596e1f2a58a2a1c8d5a05d5a74461af552b371e82f9fa4ada8342"}, + {file = "black-22.8.0-cp36-cp36m-win_amd64.whl", hash = "sha256:5594efbdc35426e35a7defa1ea1a1cb97c7dbd34c0e49af7fb593a36bd45edab"}, + {file = "black-22.8.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:a983526af1bea1e4cf6768e649990f28ee4f4137266921c2c3cee8116ae42ec3"}, + {file = "black-22.8.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3b2c25f8dea5e8444bdc6788a2f543e1fb01494e144480bc17f806178378005e"}, + {file = "black-22.8.0-cp37-cp37m-win_amd64.whl", hash = "sha256:78dd85caaab7c3153054756b9fe8c611efa63d9e7aecfa33e533060cb14b6d16"}, + {file = "black-22.8.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:cea1b2542d4e2c02c332e83150e41e3ca80dc0fb8de20df3c5e98e242156222c"}, + {file = "black-22.8.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:5b879eb439094751185d1cfdca43023bc6786bd3c60372462b6f051efa6281a5"}, + {file = "black-22.8.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:0a12e4e1353819af41df998b02c6742643cfef58282915f781d0e4dd7a200411"}, + {file = "black-22.8.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c3a73f66b6d5ba7288cd5d6dad9b4c9b43f4e8a4b789a94bf5abfb878c663eb3"}, + {file = "black-22.8.0-cp38-cp38-win_amd64.whl", hash = "sha256:e981e20ec152dfb3e77418fb616077937378b322d7b26aa1ff87717fb18b4875"}, + {file = "black-22.8.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:8ce13ffed7e66dda0da3e0b2eb1bdfc83f5812f66e09aca2b0978593ed636b6c"}, + {file = "black-22.8.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:32a4b17f644fc288c6ee2bafdf5e3b045f4eff84693ac069d87b1a347d861497"}, + {file = "black-22.8.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:0ad827325a3a634bae88ae7747db1a395d5ee02cf05d9aa7a9bd77dfb10e940c"}, + {file = "black-22.8.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:53198e28a1fb865e9fe97f88220da2e44df6da82b18833b588b1883b16bb5d41"}, + {file = "black-22.8.0-cp39-cp39-win_amd64.whl", hash = "sha256:bc4d4123830a2d190e9cc42a2e43570f82ace35c3aeb26a512a2102bce5af7ec"}, + {file = "black-22.8.0-py3-none-any.whl", hash = "sha256:d2c21d439b2baf7aa80d6dd4e3659259be64c6f49dfd0f32091063db0e006db4"}, + {file = "black-22.8.0.tar.gz", hash = "sha256:792f7eb540ba9a17e8656538701d3eb1afcb134e3b45b71f20b25c77a8db7e6e"}, ] certifi = [ {file = "certifi-2022.6.15-py3-none-any.whl", hash = "sha256:fe86415d55e84719d75f8b69414f6438ac3547d2078ab91b67e779ef69378412"}, @@ -1459,8 +1472,8 @@ packaging = [ {file = "packaging-21.3.tar.gz", hash = "sha256:dd47c42927d89ab911e606518907cc2d3a1f38bbd026385970643f9c5b8ecfeb"}, ] pathspec = [ - {file = "pathspec-0.9.0-py2.py3-none-any.whl", hash = "sha256:7d15c4ddb0b5c802d161efc417ec1a2558ea2653c2e8ad9c19098201dc1c993a"}, - {file = "pathspec-0.9.0.tar.gz", hash = "sha256:e564499435a2673d586f6b2130bb5b95f04a3ba06f81b8f895b651a3c76aabb1"}, + {file = "pathspec-0.10.0-py3-none-any.whl", hash = "sha256:aefa80ac32d5bf1f96139dca67cefb69a431beff4e6bf1168468f37d7ab87015"}, + {file = "pathspec-0.10.0.tar.gz", hash = "sha256:01eecd304ba0e6eeed188ae5fa568e99ef10265af7fd9ab737d6412b4ee0ab85"}, ] platformdirs = [ {file = "platformdirs-2.5.2-py3-none-any.whl", hash = "sha256:027d8e83a2d7de06bbac4e5ef7e023c02b863d7ea5d079477e722bb41ab25788"}, @@ -1639,8 +1652,12 @@ urllib3 = [ {file = "urllib3-1.26.12.tar.gz", hash = "sha256:3fa96cf423e6987997fc326ae8df396db2a8b7c667747d47ddd8ecba91f4a74e"}, ] virtualenv = [ - {file = "virtualenv-20.16.3-py2.py3-none-any.whl", hash = "sha256:4193b7bc8a6cd23e4eb251ac64f29b4398ab2c233531e66e40b19a6b7b0d30c1"}, - {file = "virtualenv-20.16.3.tar.gz", hash = "sha256:d86ea0bb50e06252d79e6c241507cb904fcd66090c3271381372d6221a3970f9"}, + {file = "virtualenv-20.16.4-py3-none-any.whl", hash = "sha256:035ed57acce4ac35c82c9d8802202b0e71adac011a511ff650cbcf9635006a22"}, + {file = "virtualenv-20.16.4.tar.gz", hash = "sha256:014f766e4134d0008dcaa1f95bafa0fb0f575795d07cae50b1bee514185d6782"}, +] +websocket-client = [ + {file = "websocket-client-1.4.0.tar.gz", hash = "sha256:79d730c9776f4f112f33b10b78c8d209f23b5806d9a783e296b3813fc5add2f1"}, + {file = "websocket_client-1.4.0-py3-none-any.whl", hash = "sha256:33ad3cf0aef4270b95d10a5a66b670a66be1f5ccf10ce390b3644f9eddfdca9d"}, ] yarl = [ {file = "yarl-1.8.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:abc06b97407868ef38f3d172762f4069323de52f2b70d133d096a48d72215d28"}, diff --git a/pyproject.toml b/pyproject.toml index d1b6249e..43fb9b91 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ pytest = "^7.1.2" pre-commit = "^2.20.0" nibiru-proto = "^0.14.1" shutup = "^0.2.0" +websocket-client = "^1.4.0" [tool.poetry.dev-dependencies] black = "^22.6.0" diff --git a/tests/__init__.py b/tests/__init__.py index 35799e48..e26f581b 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -2,72 +2,15 @@ import collections import logging import pprint -import sys from typing import Any, Iterable, Optional, Union import shutup -shutup.please() - - -class ColoredFormatter(logging.Formatter): - - fmt = "%(asctime)s|%(levelname)s|%(funcName)s| %(message)s" - - white = "\x1b[97;20m" - grey = "\x1b[38;20m" - green = "\x1b[32;20m" - cyan = "\x1b[36;20m" - yellow = "\x1b[33;20m" - red = "\x1b[31;20m" - bold_red = "\x1b[31;1m" - reset = "\x1b[0m" - - FORMATS = { - logging.DEBUG: fmt.format(green, reset), - logging.INFO: fmt.format(cyan, reset), - logging.WARNING: fmt.format(yellow, reset), - logging.ERROR: fmt.format(red, reset), - logging.CRITICAL: fmt.format(bold_red, reset), - } - - def format(self, record: logging.LogRecord): - """Formats a record for the logging handler. - - Args: - record (logging.LogRecord): Represents an instance of an event being - logged. - """ - log_format = self.FORMATS.get(record.levelno) - formatter = logging.Formatter(log_format, datefmt="%H:%M:%S") - return formatter.format(record=record) +from nibiru.utils import init_logger +shutup.please() -def init_test_logger() -> logging.Logger: - test_logger = logging.getLogger("test-logger") - test_logger.setLevel(logging.DEBUG) - - # Logs to stdout so we can at least see logs in GHA. - handler = logging.StreamHandler(sys.stdout) - handler.setLevel(logging.DEBUG) - - handler.setFormatter(fmt=ColoredFormatter()) - test_logger.addHandler(handler) - return test_logger - - -LOGGER: logging.Logger = init_test_logger() -"""Simple logger to use throughout the test suite. - -Examples: -```python -from tests import LOGGER -LOGGER.info("successfully executed tx staking command") -LOGGER.debug("debugging error message") -``` - -Log levels include: [debug, info, warning, error, critical] -""" +LOGGER: logging.Logger = init_logger("test-logger") def format_response(resp: Union[dict, list, str]) -> str: diff --git a/tests/conftest.py b/tests/conftest.py index 2dcc4607..a96ebaec 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -37,6 +37,7 @@ def pytest_configure(config): raise ValueError(f"Environment variable {var} is missing!") pytest.VALIDATOR_MNEMONIC = os.getenv("VALIDATOR_MNEMONIC") + pytest.WEBSOCKET_ENDPOINT = os.getenv("WEBSOCKET_ENDPOINT") pytest.ORACLE_MNEMONIC = os.getenv("ORACLE_MNEMONIC") pytest.NETWORK_INSECURE = os.getenv("NETWORK_INSECURE") != "false" diff --git a/tests/websocket_test.py b/tests/websocket_test.py new file mode 100644 index 00000000..0870d1ef --- /dev/null +++ b/tests/websocket_test.py @@ -0,0 +1,62 @@ +import time +from typing import List + +import nibiru +import nibiru.msg +from nibiru import Network, common +from nibiru.event_specs import EventCaptured +from nibiru.websocket import EventType, NibiruWebsocket +from tests import LOGGER + + +def test_websocket_listen(val_node: nibiru.Sdk, network: Network): + """ + Open a position and ensure output is correct + """ + pair = "ubtc:unusd" + + nibiru_websocket = NibiruWebsocket( + network, + [ + EventType.MarkPriceChanged, + EventType.PositionChangedEvent, + ], + ) + nibiru_websocket.start() + time.sleep(1) + + # Open a position from the validator node + LOGGER.info("Opening position") + val_node.tx.execute_msgs( + nibiru.msg.MsgOpenPosition( + sender=val_node.address, + token_pair=pair, + side=common.Side.BUY, + quote_asset_amount=10, + leverage=10, + base_asset_amount_limit=0, + ) + ) + + # Give time for events to come + LOGGER.info("Sent txs, waiting for websocket to pick it up") + time.sleep(1) + + nibiru_websocket.queue.put(None) + events: List[EventCaptured] = [] + event = 1 + while True: + event = nibiru_websocket.queue.get() + if event is None: + break + events.append(event) + + # Asserting for truth because test are running in parallel in the same chain and might result in + # duplication of markpricechanged events. + assert all( + event in [event.event_type for event in events] + for event in [ + EventType.MarkPriceChanged.value, + EventType.PositionChangedEvent.value, + ] + )