services/nodemanager/arvnodeman/test/fake_driver.py

   1 import re
   2 import urllib
   3 import ssl
   4
   5 from libcloud.compute.base import NodeSize, Node, NodeDriver, NodeState
   6 from libcloud.common.exceptions import BaseHTTPError
   7
   8 all_nodes = []
   9 create_calls = 0
  10 quota = 2
  11
  12 class FakeDriver(NodeDriver):
  13     def __init__(self, *args, **kwargs):
  14         self.name = "FakeDriver"
  15
  16     def list_sizes(self, **kwargs):
  17         return [NodeSize("Standard_D3", "Standard_D3", 3500, 200, 0, 0, self),
  18                 NodeSize("Standard_D4", "Standard_D4", 7000, 400, 0, 0, self)]
  19
  20     def list_nodes(self, **kwargs):
  21         return all_nodes
  22
  23     def create_node(self, name=None,
  24                     size=None,
  25                     image=None,
  26                     auth=None,
  27                     ex_storage_account=None,
  28                     ex_customdata=None,
  29                     ex_resource_group=None,
  30                     ex_user_name=None,
  31                     ex_tags=None,
  32                     ex_network=None):
  33         global all_nodes, create_calls
  34         create_calls += 1
  35         n = Node(name, name, NodeState.RUNNING, [], [], self, size=size, extra={"tags": ex_tags})
  36         all_nodes.append(n)
  37         ping_url = re.search(r"echo '(.*)' > /var/tmp/arv-node-data/arv-ping-url", ex_customdata).groups(1)[0] + "&instance_id=" + name
  38         ctx = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
  39         ctx.verify_mode = ssl.CERT_NONE
  40         f = urllib.urlopen(ping_url, "", context=ctx)
  41         f.close()
  42         return n
  43
  44     def destroy_node(self, cloud_node):
  45         global all_nodes
  46         all_nodes = [n for n in all_nodes if n.id != cloud_node.id]
  47         return True
  48
  49     def get_image(self, img):
  50         pass
  51
  52     def ex_create_tags(self, cloud_node, tags):
  53         pass
  54
  55 class QuotaDriver(FakeDriver):
  56     def create_node(self, name=None,
  57                     size=None,
  58                     image=None,
  59                     auth=None,
  60                     ex_storage_account=None,
  61                     ex_customdata=None,
  62                     ex_resource_group=None,
  63                     ex_user_name=None,
  64                     ex_tags=None,
  65                     ex_network=None):
  66         global all_nodes, create_calls, quota
  67         if len(all_nodes) >= quota:
  68             raise BaseHTTPError(503, "Quota exceeded")
  69         else:
  70             return super(QuotaDriver, self).create_node(name=name,
  71                     size=size,
  72                     image=image,
  73                     auth=auth,
  74                     ex_storage_account=ex_storage_account,
  75                     ex_customdata=ex_customdata,
  76                     ex_resource_group=ex_resource_group,
  77                     ex_user_name=ex_user_name,
  78                     ex_tags=ex_tags,
  79                     ex_network=ex_network)
  80
  81     def destroy_node(self, cloud_node):
  82         global all_nodes, quota
  83         all_nodes = [n for n in all_nodes if n.id != cloud_node.id]
  84         if len(all_nodes) == 0:
  85             quota = 4
  86         return True
  87
  88 class FailingDriver(FakeDriver):
  89     def create_node(self, name=None,
  90                     size=None,
  91                     image=None,
  92                     auth=None,
  93                     ex_storage_account=None,
  94                     ex_customdata=None,
  95                     ex_resource_group=None,
  96                     ex_user_name=None,
  97                     ex_tags=None,
  98                     ex_network=None):
  99         raise Exception("nope")
 100
 101 class RetryDriver(FakeDriver):
 102     def create_node(self, name=None,
 103                     size=None,
 104                     image=None,
 105                     auth=None,
 106                     ex_storage_account=None,
 107                     ex_customdata=None,
 108                     ex_resource_group=None,
 109                     ex_user_name=None,
 110                     ex_tags=None,
 111                     ex_network=None):
 112         global create_calls
 113         create_calls += 1
 114         if create_calls < 2:
 115             raise BaseHTTPError(429, "Rate limit exceeded",
 116                                 {'retry-after': '12'})
 117         else:
 118             return super(RetryDriver, self).create_node(name=name,
 119                     size=size,
 120                     image=image,
 121                     auth=auth,
 122                     ex_storage_account=ex_storage_account,
 123                     ex_customdata=ex_customdata,
 124                     ex_resource_group=ex_resource_group,
 125                     ex_user_name=ex_user_name,
 126                     ex_tags=ex_tags,
 127                     ex_network=ex_network)