antithesis_sdk.py

import os, json, random

'''
Quick wrapper around the fallback SDK
https://antithesis.com/docs/using_antithesis/sdk/fallback/assert.html
'''
class antithesis_fallback_sdk():

    def __init__(self):
        if 'ANTITHESIS_OUTPUT_DIR' not in os.environ:
            raise Exception(f"The environment variable ANTITHESIS_OUTPUT_DIR is not defined")

        self.sdk_location = f"{os.environ['ANTITHESIS_OUTPUT_DIR']}/sdk.jsonl"
        self.assertion_types = ['always', 'sometimes', 'unreachable', 'reachable']

    def _get_assertion_evaulate(self) -> dict:
        always = {
            'hit': True,
            'must_hit': True,
            'condition': True,
        }

        sometimes = {
            'hit': True,
            'must_hit': True,
            'condition': True,
        }

        reachable = {
            'hit': True,
            'must_hit': True,
            'condition': True,
        }

        unreachable = {
            'hit': True,
            'must_hit': False,
            'condition': False,
        }

        return {
            'always': always,
            'sometimes': sometimes,
            'reachable': reachable,
            'unreachable': unreachable,
        }

    def _get_assertion_declare(self):
        '''
        table of metadata for different assertions
        '''
        always = {
            'hit': False,
            'must_hit': True,
            'assert_type': 'always',
            'display_type': 'Always',
            'condition': False,
            'details': None,
        }

        sometimes = {
            'hit': False,
            'must_hit': True,
            'assert_type': 'sometimes',
            'display_type': 'Sometimes',
            'condition': False,
            'details': None,
        }

        unreachable = {
            'hit': False,
            'must_hit': False,
            'assert_type': 'reachability',
            'display_type': 'Unreachable',
            'condition': False,
            'details': None,
        }

        reachable = {
            'hit': False,
            'must_hit': True,
            'assert_type': 'reachability',
            'display_type': 'Reachable',
            'condition': False,
            'details': None,
        }

        return {
            'always': always,
            'sometimes': sometimes,
            'unreachable': unreachable,
            'reachable': reachable,
        }

    def _declare(self, assert_type:str, id: str, message:str) -> dict:
        '''
        Make the assertion declaration
        '''
        if assert_type not in self.assertion_types:
            raise Exception(f"The assertion type must be one of {', '.join(self.assertion_types)}")

        assertions_declare = self._get_assertion_declare()
        payload = {
            'antithesis_assert': assertions_declare[assert_type]
        }

        # Stubbing these out for now
        location = {
            "class": "",
            "function": "",
            "file": "",
            "begin_line": 0,
            "begin_column": 0,
        }

        payload["antithesis_assert"]["location"] = location

        return payload

    def declare(self, assert_type:str, id: str, message:str) -> None:
        payload = self._declare(assert_type, id, message)
        self._write_to_json(payload)

    def do_assert(self, assert_type:str, id: str, message:str, condition:bool, details:dict = {}) -> None:
        '''
        Write the assert evaulation to the event sink
        '''

        # Build the base payload
        payload = self._declare(assert_type, id, message)
        
        # Change the values based on evaulation
        assert_eval = self._get_assertion_evaulate()[assert_type]

        antithesis_assert = payload['antithesis_assert']
        antithesis_assert.update(assert_eval)

        payload['antithesis_assert'] = antithesis_assert

        payload['antithesis_assert']['condition'] = condition

        if bool(details):
            payload['antithesis_assert']['details'] = details

        self._write_to_json(payload)

    def always(self, declare:bool, id:str, message:str, condition:bool = False, details:dict = {}, ):
        if declare:
            self.declare('always', id, message)
        else:
            self.do_assert('always', id, message, condition, details)

    def sometimes(self, declare:bool, id:str, message:str, condition:bool = False, details:dict = {}):
        if declare:
            self.declare('sometimes', id, message)
        else:
            self.do_assert('sometimes', id, message, condition, details)

    def reachable(self, declare:bool, id:str, message:str, condition:bool = False, details:dict = {}):
        if declare:
            self.declare('reachable', id, message)
        else:
            self.do_assert('reachable', id, message, condition, details)

    def unreachable(self, declare:bool, id:str, message:str, condition:bool = False, details:dict = {}):
        if declare:
            self.declare('unreachable', id, message)
        else:
            self.do_assert('unreachable', id, message, condition, details)

    def _write_to_json(self, assertion_obj:dict) -> None:
        '''
        Appending the data to json
        '''
        with open(self.sdk_location, 'a') as fh: 
            fh.write(f'{json.dumps(assertion_obj)}\n')

    def get_random_int(self, bytes = 1) -> int:
        random_bytes = os.urandom(bytes)
        return int.from_bytes(random_bytes, byteorder='big')

    def setup_complete(self, details:dict = {}) -> None:

        message = {
            "antithesis_setup": {
                "status": "complete", 
                "details": details
            }
        }

        self._write_to_json(message)