1#!/bin/bash
2#
3# Copyright 2014 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7
8HELP="This is a script to bootstrap a localhost shard cluster for testing.\n\
9The following defaults are preconfigured but modifyable:\n\
10  SHARD_NAME: Name of the shard to register with master\n\
11  NUM_HOSTS_MASTER/SHARD: Number of hosts to add to the master/shard.\n\
12  MASTER/SHARD_BOARD: Boards to add to the master/shard\n\
13  POOL: Pool to use for the hosts."
14
15
16# Invalidate (delete) the hosts/labels/shard instead of adding them.
17# Typically used to refresh a botched cluster.
18INVALIDATE_ALL=0
19AT_DIR=/usr/local/autotest
20
21# See VagrantFile for details on how these afes are setup.
22AFE=localhost:8001
23SHARD_NAME=localhost:8004
24
25# Number of hosts on master and shard. They will
26# get autoassigned generic names like test_hostX.
27NUM_HOSTS_MASTER=10
28NUM_HOSTS_SHARD=5
29NUM_FREON_HOSTS_SHARD=5
30
31# A host can only have a single board. Jobs are sent
32# to the shard based on the board.
33MASTER_BOARD=board:link
34SHARD_BOARD=board:stumpy
35SHARD_FREON_BOARD=board:stumpy_freon
36
37# All hosts need to be in a pool.
38POOL=pool:bot
39
40y_n_prompt() {
41  read -r -p "Are you sure? [y/N] " response
42  if [[ $response =~ ^([yY][eE][sS]|[yY])$ ]]; then
43    return 0
44  else
45    return 1
46  fi
47}
48
49while getopts ":h" opt; do
50  case $opt in
51    h)
52      echo -e "${HELP}" >&2
53      exit 0
54      ;;
55  esac
56done
57
58atest_hosts() {
59  hosts=("${!1}")
60  labels="${2}"
61  hostnames=''
62  for H in ${hosts[*]}; do
63    if [ "$hostnames" ]; then
64      hostnames="$hostnames,$H"
65    else
66      hostnames=$H
67    fi
68  done
69  if [ $INVALIDATE_ALL -eq 1 ]; then
70    $AT_DIR/cli/atest host delete $hostnames --web $AFE
71    $AT_DIR/cli/atest label delete $labels --web $AFE
72  else
73    $AT_DIR/cli/atest host create $hostnames --web $AFE
74    $AT_DIR/cli/atest label add -m $hostnames $labels --web $AFE
75  fi
76}
77
78MASTER_HOSTS=()
79s=1
80e=$NUM_HOSTS_MASTER
81for i in $(seq $s $e); do
82  MASTER_HOSTS[$i]=test_host$i;
83done
84
85SHARD_HOSTS=()
86s=$(($e+1))
87e=$(($NUM_HOSTS_SHARD+$e))
88for i in $(seq $s $e); do
89  SHARD_HOSTS[$i]=test_host$i;
90done
91
92SHARD_FREON_HOSTS=()
93s=$(($e+1))
94e=$(($NUM_FREON_HOSTS_SHARD+$e))
95for i in $(seq $s $e); do
96  SHARD_FREON_HOSTS[$i]=test_host$i;
97done
98
99operation='Adding: '
100if [ $INVALIDATE_ALL -eq 1 ]; then
101  operation='Removing '
102fi
103
104printf '%s following hosts to master \n' $operation
105echo ${MASTER_HOSTS[*]}
106if $(y_n_prompt); then
107  atest_hosts MASTER_HOSTS[*] $POOL,$MASTER_BOARD
108fi
109
110printf '\n\n%s following hosts to shard \n' $operation
111echo ${SHARD_HOSTS[*]}
112if $(y_n_prompt); then
113  atest_hosts SHARD_HOSTS[*] $POOL,$SHARD_BOARD
114fi
115
116printf '\n\n%s following hosts to shard \n' $operation
117echo ${SHARD_FREON_HOSTS[*]}
118if $(y_n_prompt); then
119    atest_hosts SHARD_FREON_HOSTS[*] $POOL,$SHARD_FREON_BOARD
120fi
121
122printf '\n\n%s shard \n' $operation
123echo $SHARD_NAME
124if $(y_n_prompt); then
125  if [ $INVALIDATE_ALL -eq 1 ]; then
126    $AT_DIR/cli/atest shard delete $SHARD_NAME --web $AFE
127  else
128    $AT_DIR/cli/atest shard create $SHARD_NAME -l $SHARD_BOARD --web $AFE
129  fi
130fi
131