import libvirt
import threading
import time
import sys
import optparse

poolname = "default"
volname = "nonsparsetest"
testvol = "testvol"

uri = "qemu:///system"

testvolxml = """
<volume>
  <name>%s</name>
    <capacity>1048576000</capacity>
    <allocation>0</allocation>
    <target>
        <format type='raw'/>
    </target>
</volume>
""" % testvol

volxml = """
<volume>
  <name>%s</name>
    <capacity>1048576000</capacity>
    <allocation>800000000</allocation>
    <target>
        <format type='raw'/>
    </target>
</volume>
""" % volname

failvol = """
<volume>
  <name>%s</name>
    <capacity>1048576000</capacity>
    <allocation>1048576000</allocation>
    <target>
        <format type='bochs'/>
    </target>
</volume>
""" % volname

poolxml = """
<pool type='dir'>
  <name>default</name>
  <target>
    <path>/var/lib/libvirt/images</path>
  </target>
</pool>
"""

# Helper functions

def exception_wrapper(cmd, args):
    try:
        cmd(*args)
    except Exception, e:
        print str(e)

def make_test_vol():
    pool = get_pool()
    pool.createXML(testvolxml, 0)

def del_vol(name=volname):
    pool = get_pool()
    vol = pool.storageVolLookupByName(name)
    vol.delete(0)

def del_pool():
    pool = get_pool()
    pool.destroy()
    pool.undefine()

def define_pool():
    conn = libvirt.open(uri)
    pool = conn.storagePoolDefineXML(poolxml, 0)
    pool.create(0)
    pool.setAutostart(True)

def allocate_thread(xml=volxml):
    try:
        del_vol()
    except:
        pass

    pool = get_pool()
    print "creating vol in thread"
    vol = pool.createXML(xml, 0)
    print "creating vol complete."

def info_thread(vol):

    for i in range(0, 40):
        time.sleep(.5)
        print vol.info()

def get_pool(name=poolname):
    conn = libvirt.open(uri)
    return conn.storagePoolLookupByName(name)

def get_vol(name=volname):
    pool = get_pool()
    return pool.storageVolLookupByName(name)

def cmd_vol_create(xml=volxml):

    exception_wrapper(define_pool, ())

    pool = get_pool()
    print pool.listVolumes()

    t = threading.Thread(target=allocate_thread, name="Allocating",
                         args=(xml,))
    t.start()

    time.sleep(5)
    print "\nRefreshing pool and dumping list"
    pool.refresh(0)
    print pool.listVolumes()

def cmd_vol_fail():
    cmd_vol_create(failvol)

def cmd_vol_poll():
    cmd_vol_create()
    vol = get_vol()

    t = threading.Thread(target=info_thread, name="Getting info",
                         args=(vol,))
    t.start()

def main():
    import optparse

    parser = optparse.OptionParser()

    parser.add_option("", "--vol-create", action="store_true",
                      dest="vol_create",
                      help="Create a volume nonsparse volume that should "
                           "succeed")
    parser.add_option("", "--vol-create-info", action="store_true",
                      dest="vol_create_info",
                      help="Create a volume nonsparse volume that should "
                           "succeed, and list info in a loop")
    parser.add_option("", "--vol-create-fail", action="store_true",
                      dest="vol_fail",
                      help="Create a volume that will fail at the allocate "
                           "stage")

    options, ignore = parser.parse_args()

    if options.vol_create:
        cmd_vol_create()
    elif options.vol_create_info:
        cmd_vol_poll()
    elif options.vol_fail:
        cmd_vol_fail()
    else:
        parser.print_help()
        sys.exit(1)

if __name__ == "__main__":
    main()
