import os, sys
sys.path.append(os.path.split(os.path.realpath(__file__))[0] + "/../lib")
from cnstream import *
from cnstream_cpptest import *

class CustomSourceModule(SourceModule):
    def __init__(self, name):
        SourceModule.__init__(self, name)

    def open(self, params):
        return True

    def close(self):
        return 

class CustomSourceHandler(SourceHandler):
    def __init__(self, source_module, name):
        SourceHandler.__init__(self, source_module, name)
    
    def open(self):
        return True

    def close(self):
        print("handler close")
        return


def test_source():
    source = CustomSourceModule("source_test")
    handler0 = CustomSourceHandler(source, "stream_id_0")
    handler1 = CustomSourceHandler(source, "stream_id_1")
    handler2 = CustomSourceHandler(source, "stream_id_2")
    assert handler0.open()
    assert "stream_id_0" == handler0.get_stream_id()
    data = handler0.create_frame_info()
    assert handler0.send_data(data)

    assert 0 == source.add_source(handler0)
    assert 0 == source.add_source(handler1)
    assert 0 == source.add_source(handler2)
    assert handler0 == source.get_source_handler("stream_id_0")
    assert 0 == source.remove_source("stream_id_0")
    assert None == source.get_source_handler("stream_id_0")
    assert handler1 == source.get_source_handler("stream_id_1")
    assert 0 == source.remove_source(handler1)
    assert None == source.get_source_handler("stream_id_1")
    assert handler2 == source.get_source_handler("stream_id_2")
    assert 0 == source.remove_sources()
    assert None == source.get_source_handler("stream_id_2")