add whole new app for balance prediction, replace uwsgi, update everything
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
FROM python:3.9.0
|
||||
FROM python:3.9.1
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
WORKDIR /code/
|
||||
COPY requirements.txt .
|
||||
|
||||
0
core/__init__.py
Normal file
0
core/__init__.py
Normal file
15
core/admin.py
Normal file
15
core/admin.py
Normal file
@@ -0,0 +1,15 @@
|
||||
from django.contrib import admin
|
||||
|
||||
from core.models import Subject, Transaction
|
||||
from financeplanner.admin import admin_site
|
||||
|
||||
|
||||
@admin.register(Subject, site=admin_site)
|
||||
class SubjectAdmin(admin.ModelAdmin):
|
||||
pass
|
||||
|
||||
|
||||
@admin.register(Transaction, site=admin_site)
|
||||
class TransactionAdmin(admin.ModelAdmin):
|
||||
list_display = ("subject", "amount", "booking_date",)
|
||||
ordering = ("-booking_date",)
|
||||
5
core/apps.py
Normal file
5
core/apps.py
Normal file
@@ -0,0 +1,5 @@
|
||||
from django.apps import AppConfig
|
||||
|
||||
|
||||
class CoreConfig(AppConfig):
|
||||
name = 'core'
|
||||
0
core/management/__init__.py
Normal file
0
core/management/__init__.py
Normal file
0
core/management/commands/__init__.py
Normal file
0
core/management/commands/__init__.py
Normal file
74
core/management/commands/predict_balance.py
Normal file
74
core/management/commands/predict_balance.py
Normal file
@@ -0,0 +1,74 @@
|
||||
from datetime import datetime
|
||||
from decimal import Decimal
|
||||
|
||||
from dateutil.relativedelta import relativedelta
|
||||
from django.core.management import BaseCommand
|
||||
|
||||
from core.models import Subject
|
||||
from core.prediction import predict_transactions
|
||||
|
||||
past_lookup_delta = relativedelta(weeks=2)
|
||||
future_lookup_delta = relativedelta(months=2)
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument(
|
||||
"--start", default="0", help="the current balance to use as a starting point for prediction",
|
||||
)
|
||||
|
||||
def handle(self, *args, **options):
|
||||
start_balance = Decimal(options["start"])
|
||||
|
||||
today = datetime.now().date()
|
||||
past_lookup_bound = today - past_lookup_delta
|
||||
future_lookup_bound = today + future_lookup_delta
|
||||
|
||||
transaction_dict = {}
|
||||
for subject in Subject.objects.all():
|
||||
transactions = []
|
||||
for tr in subject.transactions.order_by("booking_date"):
|
||||
if past_lookup_bound <= tr.booking_date <= future_lookup_bound:
|
||||
transactions.append(tr)
|
||||
predicted_transaction, prediction_info = predict_transactions(subject)
|
||||
if predicted_transaction:
|
||||
first_predicted_date = predicted_transaction[0].booking_date
|
||||
if first_predicted_date >= past_lookup_bound:
|
||||
# if two weeks after the first predicted transaction have passed, the subject is considered done
|
||||
for tr in predicted_transaction:
|
||||
if past_lookup_bound <= tr.booking_date <= future_lookup_bound:
|
||||
transactions.append(tr)
|
||||
transaction_dict[subject] = (transactions, prediction_info)
|
||||
|
||||
future_transactions = []
|
||||
for subject, prediction in transaction_dict.items():
|
||||
transactions, info = prediction[0], prediction[1]
|
||||
print(f">>> {subject}")
|
||||
if info:
|
||||
rec_days, rec_months, day_of_month = \
|
||||
info["recurring_days"], info["recurring_months"], info["day_of_month"]
|
||||
if rec_months and day_of_month:
|
||||
print(f"~~~ predicted transaction on day-of-month {day_of_month} every {rec_months} month(s)")
|
||||
elif rec_months:
|
||||
print(f"~~~ predicted transaction every {rec_months} month(s)")
|
||||
elif rec_days:
|
||||
print(f"~~~ predicted transaction every {rec_days} day(s)")
|
||||
else:
|
||||
print("~~~ no prediction possible")
|
||||
for tr in transactions:
|
||||
print(f" {tr.booking_date:%d.%m.%Y} | {tr.amount} € | {'stored' if tr.pk else 'predicted'}")
|
||||
if tr.booking_date > today:
|
||||
future_transactions.append(tr)
|
||||
print()
|
||||
|
||||
current_balance = start_balance
|
||||
print(f"starting calculation with amount {current_balance} €")
|
||||
for tr in sorted(future_transactions, key=lambda t: t.booking_date):
|
||||
current_balance += tr.amount
|
||||
print(
|
||||
f"{tr.booking_date:%d.%m.%Y} | "
|
||||
f"{tr.subject.name:<18} | "
|
||||
f"{tr.amount:>8} € ==> "
|
||||
f"{current_balance:>8} €"
|
||||
)
|
||||
48
core/migrations/0001_initial.py
Normal file
48
core/migrations/0001_initial.py
Normal file
@@ -0,0 +1,48 @@
|
||||
# Generated by Django 3.1.5 on 2021-01-04 13:41
|
||||
|
||||
import core.models
|
||||
from django.db import migrations, models
|
||||
import django.db.models.deletion
|
||||
import django.utils.timezone
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
initial = True
|
||||
|
||||
dependencies = [
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.CreateModel(
|
||||
name='Balance',
|
||||
fields=[
|
||||
('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
|
||||
('amount', core.models.PriceField(decimal_places=2, max_digits=7)),
|
||||
('date', models.DateField()),
|
||||
],
|
||||
options={
|
||||
'get_latest_by': 'date',
|
||||
},
|
||||
),
|
||||
migrations.CreateModel(
|
||||
name='Subject',
|
||||
fields=[
|
||||
('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
|
||||
('name', models.CharField(max_length=64)),
|
||||
('created_time', models.DateTimeField(default=django.utils.timezone.now)),
|
||||
],
|
||||
),
|
||||
migrations.CreateModel(
|
||||
name='Transaction',
|
||||
fields=[
|
||||
('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
|
||||
('amount', core.models.PriceField(decimal_places=2, max_digits=7)),
|
||||
('booking_date', models.DateField()),
|
||||
('subject', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='transactions', related_query_name='transaction', to='core.subject')),
|
||||
],
|
||||
options={
|
||||
'get_latest_by': 'booking_date',
|
||||
},
|
||||
),
|
||||
]
|
||||
0
core/migrations/__init__.py
Normal file
0
core/migrations/__init__.py
Normal file
49
core/models.py
Normal file
49
core/models.py
Normal file
@@ -0,0 +1,49 @@
|
||||
from decimal import Decimal
|
||||
|
||||
from django.db import models
|
||||
from django.utils import timezone
|
||||
|
||||
|
||||
class PriceField(models.DecimalField):
|
||||
# allowing numbers until 99,999.99
|
||||
max_digits = 7
|
||||
decimal_places = 2
|
||||
|
||||
def to_python(self, value):
|
||||
value = super().to_python(value)
|
||||
if isinstance(value, Decimal):
|
||||
return value.quantize(Decimal(10) ** -self.decimal_places)
|
||||
return value
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['max_digits'] = self.max_digits
|
||||
kwargs['decimal_places'] = self.decimal_places
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
class Subject(models.Model):
|
||||
name = models.CharField(max_length=64)
|
||||
created_time = models.DateTimeField(default=timezone.now)
|
||||
|
||||
def __str__(self):
|
||||
return self.name
|
||||
|
||||
|
||||
class Transaction(models.Model):
|
||||
class Meta:
|
||||
get_latest_by = "booking_date"
|
||||
|
||||
amount = PriceField()
|
||||
booking_date = models.DateField()
|
||||
subject = models.ForeignKey(
|
||||
to=Subject, on_delete=models.CASCADE,
|
||||
related_name="transactions", related_query_name="transaction",
|
||||
)
|
||||
|
||||
|
||||
class Balance(models.Model):
|
||||
class Meta:
|
||||
get_latest_by = "date"
|
||||
|
||||
amount = PriceField()
|
||||
date = models.DateField()
|
||||
87
core/prediction.py
Normal file
87
core/prediction.py
Normal file
@@ -0,0 +1,87 @@
|
||||
from decimal import Decimal
|
||||
|
||||
from dateutil.relativedelta import relativedelta
|
||||
|
||||
from core.models import Subject, Transaction
|
||||
|
||||
days_per_year = Decimal("365.25")
|
||||
days_per_month = days_per_year / Decimal("12")
|
||||
|
||||
monthly_distance_threshold = 3
|
||||
certain_day_of_month_threshold = 3
|
||||
|
||||
|
||||
def _group_follow_up_objects(objects):
|
||||
"""
|
||||
Creates a list of each follow-up pair of objects within a given list, like this:
|
||||
[1, 2, 3] -> [(1, 2), (2, 3)]
|
||||
"""
|
||||
tuples = []
|
||||
if len(objects) >= 2:
|
||||
for idx in range(len(objects) - 1):
|
||||
tuples.append((objects[idx], objects[idx + 1]))
|
||||
return tuples
|
||||
|
||||
|
||||
def predict_amount(subject: Subject):
|
||||
if subject.transactions.exists():
|
||||
return subject.transactions.latest().amount
|
||||
return None
|
||||
|
||||
|
||||
def predict_booking_dates(subject: Subject):
|
||||
dates, recurring_days, recurring_months, day_of_month = [], None, None, None
|
||||
if subject.transactions.count() < 2:
|
||||
return dates, recurring_days, recurring_months, day_of_month
|
||||
|
||||
existing_transactions = subject.transactions.order_by("booking_date")
|
||||
last_date = existing_transactions.latest().booking_date
|
||||
one_year_later = last_date + relativedelta(years=1)
|
||||
|
||||
transaction_tuples = _group_follow_up_objects(existing_transactions)
|
||||
date_deltas = [second.booking_date - first.booking_date for first, second in transaction_tuples]
|
||||
average_delta_in_days = Decimal(sum(delta.days for delta in date_deltas)) / Decimal(len(date_deltas))
|
||||
average_delta_in_months = average_delta_in_days / days_per_month
|
||||
|
||||
days_per_month_mod = average_delta_in_days % days_per_month
|
||||
distance_from_days_per_month = min(days_per_month_mod, days_per_month - days_per_month_mod)
|
||||
if distance_from_days_per_month <= monthly_distance_threshold:
|
||||
# transactions can be considered to happen every n months
|
||||
recurring_months = round(average_delta_in_months)
|
||||
while last_date < one_year_later and len(dates) < 10:
|
||||
last_date += relativedelta(months=recurring_months)
|
||||
dates.append(last_date)
|
||||
|
||||
days_of_month = [t.booking_date.day for t in existing_transactions]
|
||||
if max(days_of_month) - min(days_of_month) <= certain_day_of_month_threshold:
|
||||
# since transactions occurred in a close range around a certain day of month, we can
|
||||
# improve our prediction
|
||||
day_of_month = round(sum(days_of_month) / len(days_of_month))
|
||||
for idx in range(len(dates)):
|
||||
dates[idx] = dates[idx].replace(day=day_of_month)
|
||||
|
||||
else:
|
||||
# there is no monthly pattern, just add the average delta to determine new dates
|
||||
recurring_days = round(average_delta_in_days)
|
||||
while last_date < one_year_later and len(dates) < 10:
|
||||
last_date += relativedelta(days=recurring_days)
|
||||
dates.append(last_date)
|
||||
|
||||
return dates, recurring_days, recurring_months, day_of_month
|
||||
|
||||
|
||||
def predict_transactions(subject: Subject):
|
||||
"""
|
||||
Analyze existing transactions of a given subject and predict future transactions, up to one year in advance.
|
||||
"""
|
||||
transactions, prediction_info = [], None
|
||||
amount = predict_amount(subject)
|
||||
dates, rec_days, rec_months, day_of_month = predict_booking_dates(subject)
|
||||
if amount and dates:
|
||||
transactions = [Transaction(amount=amount, booking_date=date, subject=subject) for date in dates]
|
||||
prediction_info = {
|
||||
"recurring_days": rec_days,
|
||||
"recurring_months": rec_months,
|
||||
"day_of_month": day_of_month,
|
||||
}
|
||||
return transactions, prediction_info
|
||||
0
core/tests/__init__.py
Normal file
0
core/tests/__init__.py
Normal file
114
core/tests/test_prediction.py
Normal file
114
core/tests/test_prediction.py
Normal file
@@ -0,0 +1,114 @@
|
||||
from datetime import date
|
||||
|
||||
from django.test import TestCase
|
||||
|
||||
from core.models import Subject, Transaction
|
||||
from core.prediction import predict_transactions
|
||||
|
||||
|
||||
class PredictionTestCase(TestCase):
|
||||
|
||||
def test_not_enough_data(self):
|
||||
subject = Subject.objects.create(name="rent")
|
||||
predicted_transactions, prediction_info = predict_transactions(subject)
|
||||
self.assertEqual(predicted_transactions, [])
|
||||
self.assertIsNone(prediction_info)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 6, 2), subject=subject)
|
||||
predicted_transactions, prediction_info = predict_transactions(subject)
|
||||
self.assertEqual(predicted_transactions, [])
|
||||
self.assertIsNone(prediction_info)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 7, 2), subject=subject)
|
||||
predicted_transactions, prediction_info = predict_transactions(subject)
|
||||
self.assertNotEqual(predicted_transactions, [])
|
||||
self.assertIsNotNone(prediction_info)
|
||||
|
||||
def test_every_month(self):
|
||||
subject = Subject.objects.create(name="rent")
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 6, 2), subject=subject)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 7, 2), subject=subject)
|
||||
predicted_transactions, prediction_info = predict_transactions(subject)
|
||||
first, second, last = predicted_transactions[0], predicted_transactions[1], predicted_transactions[-1]
|
||||
self.assertEqual(first.booking_date, date(2020, 8, 2))
|
||||
self.assertEqual(second.booking_date, date(2020, 9, 2))
|
||||
self.assertEqual(last.booking_date, date(2021, 5, 2))
|
||||
for tr in predicted_transactions:
|
||||
self.assertEqual(tr.amount, -333)
|
||||
self.assertDictEqual(
|
||||
prediction_info,
|
||||
{"recurring_days": None, "recurring_months": 1, "day_of_month": 2},
|
||||
)
|
||||
|
||||
def test_every_3_months(self):
|
||||
subject = Subject.objects.create(name="rent")
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 6, 2), subject=subject)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 9, 2), subject=subject)
|
||||
predicted_transactions, prediction_info = predict_transactions(subject)
|
||||
first, second, last = predicted_transactions[0], predicted_transactions[1], predicted_transactions[-1]
|
||||
self.assertEqual(first.booking_date, date(2020, 12, 2))
|
||||
self.assertEqual(second.booking_date, date(2021, 3, 2))
|
||||
self.assertEqual(last.booking_date, date(2021, 9, 2))
|
||||
for tr in predicted_transactions:
|
||||
self.assertEqual(tr.amount, -333)
|
||||
self.assertDictEqual(
|
||||
prediction_info,
|
||||
{"recurring_days": None, "recurring_months": 3, "day_of_month": 2},
|
||||
)
|
||||
|
||||
def test_every_year(self):
|
||||
subject = Subject.objects.create(name="rent")
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2019, 6, 2), subject=subject)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 6, 2), subject=subject)
|
||||
predicted_transactions, prediction_info = predict_transactions(subject)
|
||||
self.assertEqual(len(predicted_transactions), 1)
|
||||
trans = predicted_transactions[0]
|
||||
self.assertEqual(trans.booking_date, date(2021, 6, 2))
|
||||
self.assertEqual(trans.amount, -333)
|
||||
self.assertDictEqual(
|
||||
prediction_info,
|
||||
{"recurring_days": None, "recurring_months": 12, "day_of_month": 2},
|
||||
)
|
||||
|
||||
def test_monthly_varying_begin_end(self):
|
||||
subject = Subject.objects.create(name="rent")
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 1, 31), subject=subject)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 3, 1), subject=subject)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 4, 1), subject=subject)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 4, 30), subject=subject)
|
||||
predicted_transactions, prediction_info = predict_transactions(subject)
|
||||
first, second, last = predicted_transactions[0], predicted_transactions[1], predicted_transactions[-1]
|
||||
self.assertEqual(first.booking_date, date(2020, 5, 30))
|
||||
self.assertEqual(second.booking_date, date(2020, 6, 30))
|
||||
self.assertEqual(last.booking_date, date(2021, 2, 28))
|
||||
for tr in predicted_transactions:
|
||||
self.assertEqual(tr.amount, -333)
|
||||
self.assertDictEqual(
|
||||
prediction_info,
|
||||
{"recurring_days": None, "recurring_months": 1, "day_of_month": None},
|
||||
)
|
||||
|
||||
def test_no_monthly_pattern(self):
|
||||
subject = Subject.objects.create(name="rent")
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 1, 15), subject=subject)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 2, 10), subject=subject)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 4, 25), subject=subject)
|
||||
Transaction.objects.create(amount=-333, booking_date=date(2020, 5, 5), subject=subject)
|
||||
predicted_transactions, prediction_info = predict_transactions(subject)
|
||||
first, second, last = predicted_transactions[0], predicted_transactions[1], predicted_transactions[-1]
|
||||
self.assertEqual(first.booking_date, date(2020, 6, 11))
|
||||
self.assertEqual(second.booking_date, date(2020, 7, 18))
|
||||
self.assertEqual(last.booking_date, date(2021, 5, 10))
|
||||
for tr in predicted_transactions:
|
||||
self.assertEqual(tr.amount, -333)
|
||||
self.assertDictEqual(
|
||||
prediction_info,
|
||||
{"recurring_days": 37, "recurring_months": None, "day_of_month": None},
|
||||
)
|
||||
|
||||
def test_amount_change(self):
|
||||
subject = Subject.objects.create(name="rent")
|
||||
Transaction.objects.create(amount=1337, booking_date=date(2020, 6, 2), subject=subject)
|
||||
Transaction.objects.create(amount=1312, booking_date=date(2020, 7, 2), subject=subject)
|
||||
Transaction.objects.create(amount=404, booking_date=date(2020, 8, 2), subject=subject)
|
||||
predicted_transactions, prediction_info = predict_transactions(subject)
|
||||
for tr in predicted_transactions:
|
||||
self.assertEqual(tr.amount, 404)
|
||||
47
core/tests/test_price_field.py
Normal file
47
core/tests/test_price_field.py
Normal file
@@ -0,0 +1,47 @@
|
||||
from decimal import Decimal, InvalidOperation
|
||||
|
||||
from django.test import TestCase
|
||||
from django.utils import timezone
|
||||
|
||||
from core.models import Balance
|
||||
|
||||
|
||||
def _create_and_reload_balance(amount):
|
||||
bal = Balance.objects.create(amount=amount, date=timezone.now().date())
|
||||
bal.refresh_from_db()
|
||||
return bal
|
||||
|
||||
|
||||
class PriceFieldTestCase(TestCase):
|
||||
|
||||
def test_normal_decimal(self):
|
||||
bal = _create_and_reload_balance(Decimal("13.37"))
|
||||
self.assertEqual(bal.amount, Decimal("13.37"))
|
||||
|
||||
def test_too_big_decimal(self):
|
||||
with self.assertRaises(InvalidOperation):
|
||||
_create_and_reload_balance(Decimal("13371337.00"))
|
||||
|
||||
def test_rounding_decimal(self):
|
||||
bal = _create_and_reload_balance(Decimal("1.337"))
|
||||
self.assertEqual(bal.amount, Decimal("1.34"))
|
||||
|
||||
def test_normal_int(self):
|
||||
bal = _create_and_reload_balance(1337)
|
||||
self.assertEqual(bal.amount, Decimal("1337.00"))
|
||||
|
||||
def test_too_big_int(self):
|
||||
with self.assertRaises(InvalidOperation):
|
||||
_create_and_reload_balance(13371337)
|
||||
|
||||
def test_normal_float(self):
|
||||
bal = _create_and_reload_balance(13.37)
|
||||
self.assertEqual(bal.amount, Decimal("13.37"))
|
||||
|
||||
def test_too_big_float(self):
|
||||
with self.assertRaises(InvalidOperation):
|
||||
_create_and_reload_balance(13371337.00)
|
||||
|
||||
def test_rounding_float(self):
|
||||
bal = _create_and_reload_balance(1.337)
|
||||
self.assertEqual(bal.amount, Decimal("1.34"))
|
||||
3
core/views.py
Normal file
3
core/views.py
Normal file
@@ -0,0 +1,3 @@
|
||||
from django.shortcuts import render
|
||||
|
||||
# Create your views here.
|
||||
@@ -3,12 +3,10 @@ set -e
|
||||
|
||||
if [ "$DJANGO_PRODUCTION_MODE" == "true" ]; then
|
||||
echo "starting production server ..."
|
||||
./wait-for-it.sh --host=${DB_HOST} --port=${DB_PORT} --strict --timeout=20 -- \
|
||||
uwsgi --ini uwsgi.ini
|
||||
gunicorn --bind=0.0.0.0:8000 --workers=2 financeproject.wsgi
|
||||
|
||||
else
|
||||
echo "starting development server ..."
|
||||
./wait-for-it.sh --host=${DB_HOST} --port=${DB_PORT} --strict --timeout=20 -- \
|
||||
python manage.py runserver 0.0.0.0:8000
|
||||
python manage.py runserver 0.0.0.0:8000
|
||||
|
||||
fi
|
||||
|
||||
@@ -44,6 +44,7 @@ else:
|
||||
|
||||
INSTALLED_APPS = [
|
||||
"financeplanner",
|
||||
"core",
|
||||
"django.contrib.admin",
|
||||
"django.contrib.auth",
|
||||
"django.contrib.contenttypes",
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
django==3.1.4
|
||||
django==3.1.5
|
||||
psycopg2-binary==2.8.6
|
||||
uwsgi==2.0.19.1
|
||||
gunicorn==20.0.4
|
||||
python-dateutil==2.8.1
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
[uwsgi]
|
||||
chdir = /code
|
||||
module = financeproject.wsgi:application
|
||||
http = 0.0.0.0:8000
|
||||
workers = 2
|
||||
master = true
|
||||
178
wait-for-it.sh
178
wait-for-it.sh
@@ -1,178 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
# Use this script to test if a given TCP host/port are available
|
||||
|
||||
WAITFORIT_cmdname=${0##*/}
|
||||
|
||||
echoerr() { if [[ $WAITFORIT_QUIET -ne 1 ]]; then echo "$@" 1>&2; fi }
|
||||
|
||||
usage()
|
||||
{
|
||||
cat << USAGE >&2
|
||||
Usage:
|
||||
$WAITFORIT_cmdname host:port [-s] [-t timeout] [-- command args]
|
||||
-h HOST | --host=HOST Host or IP under test
|
||||
-p PORT | --port=PORT TCP port under test
|
||||
Alternatively, you specify the host and port as host:port
|
||||
-s | --strict Only execute subcommand if the test succeeds
|
||||
-q | --quiet Don't output any status messages
|
||||
-t TIMEOUT | --timeout=TIMEOUT
|
||||
Timeout in seconds, zero for no timeout
|
||||
-- COMMAND ARGS Execute command with args after the test finishes
|
||||
USAGE
|
||||
exit 1
|
||||
}
|
||||
|
||||
wait_for()
|
||||
{
|
||||
if [[ $WAITFORIT_TIMEOUT -gt 0 ]]; then
|
||||
echoerr "$WAITFORIT_cmdname: waiting $WAITFORIT_TIMEOUT seconds for $WAITFORIT_HOST:$WAITFORIT_PORT"
|
||||
else
|
||||
echoerr "$WAITFORIT_cmdname: waiting for $WAITFORIT_HOST:$WAITFORIT_PORT without a timeout"
|
||||
fi
|
||||
WAITFORIT_start_ts=$(date +%s)
|
||||
while :
|
||||
do
|
||||
if [[ $WAITFORIT_ISBUSY -eq 1 ]]; then
|
||||
nc -z $WAITFORIT_HOST $WAITFORIT_PORT
|
||||
WAITFORIT_result=$?
|
||||
else
|
||||
(echo > /dev/tcp/$WAITFORIT_HOST/$WAITFORIT_PORT) >/dev/null 2>&1
|
||||
WAITFORIT_result=$?
|
||||
fi
|
||||
if [[ $WAITFORIT_result -eq 0 ]]; then
|
||||
WAITFORIT_end_ts=$(date +%s)
|
||||
echoerr "$WAITFORIT_cmdname: $WAITFORIT_HOST:$WAITFORIT_PORT is available after $((WAITFORIT_end_ts - WAITFORIT_start_ts)) seconds"
|
||||
break
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
return $WAITFORIT_result
|
||||
}
|
||||
|
||||
wait_for_wrapper()
|
||||
{
|
||||
# In order to support SIGINT during timeout: http://unix.stackexchange.com/a/57692
|
||||
if [[ $WAITFORIT_QUIET -eq 1 ]]; then
|
||||
timeout $WAITFORIT_BUSYTIMEFLAG $WAITFORIT_TIMEOUT $0 --quiet --child --host=$WAITFORIT_HOST --port=$WAITFORIT_PORT --timeout=$WAITFORIT_TIMEOUT &
|
||||
else
|
||||
timeout $WAITFORIT_BUSYTIMEFLAG $WAITFORIT_TIMEOUT $0 --child --host=$WAITFORIT_HOST --port=$WAITFORIT_PORT --timeout=$WAITFORIT_TIMEOUT &
|
||||
fi
|
||||
WAITFORIT_PID=$!
|
||||
trap "kill -INT -$WAITFORIT_PID" INT
|
||||
wait $WAITFORIT_PID
|
||||
WAITFORIT_RESULT=$?
|
||||
if [[ $WAITFORIT_RESULT -ne 0 ]]; then
|
||||
echoerr "$WAITFORIT_cmdname: timeout occurred after waiting $WAITFORIT_TIMEOUT seconds for $WAITFORIT_HOST:$WAITFORIT_PORT"
|
||||
fi
|
||||
return $WAITFORIT_RESULT
|
||||
}
|
||||
|
||||
# process arguments
|
||||
while [[ $# -gt 0 ]]
|
||||
do
|
||||
case "$1" in
|
||||
*:* )
|
||||
WAITFORIT_hostport=(${1//:/ })
|
||||
WAITFORIT_HOST=${WAITFORIT_hostport[0]}
|
||||
WAITFORIT_PORT=${WAITFORIT_hostport[1]}
|
||||
shift 1
|
||||
;;
|
||||
--child)
|
||||
WAITFORIT_CHILD=1
|
||||
shift 1
|
||||
;;
|
||||
-q | --quiet)
|
||||
WAITFORIT_QUIET=1
|
||||
shift 1
|
||||
;;
|
||||
-s | --strict)
|
||||
WAITFORIT_STRICT=1
|
||||
shift 1
|
||||
;;
|
||||
-h)
|
||||
WAITFORIT_HOST="$2"
|
||||
if [[ $WAITFORIT_HOST == "" ]]; then break; fi
|
||||
shift 2
|
||||
;;
|
||||
--host=*)
|
||||
WAITFORIT_HOST="${1#*=}"
|
||||
shift 1
|
||||
;;
|
||||
-p)
|
||||
WAITFORIT_PORT="$2"
|
||||
if [[ $WAITFORIT_PORT == "" ]]; then break; fi
|
||||
shift 2
|
||||
;;
|
||||
--port=*)
|
||||
WAITFORIT_PORT="${1#*=}"
|
||||
shift 1
|
||||
;;
|
||||
-t)
|
||||
WAITFORIT_TIMEOUT="$2"
|
||||
if [[ $WAITFORIT_TIMEOUT == "" ]]; then break; fi
|
||||
shift 2
|
||||
;;
|
||||
--timeout=*)
|
||||
WAITFORIT_TIMEOUT="${1#*=}"
|
||||
shift 1
|
||||
;;
|
||||
--)
|
||||
shift
|
||||
WAITFORIT_CLI=("$@")
|
||||
break
|
||||
;;
|
||||
--help)
|
||||
usage
|
||||
;;
|
||||
*)
|
||||
echoerr "Unknown argument: $1"
|
||||
usage
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [[ "$WAITFORIT_HOST" == "" || "$WAITFORIT_PORT" == "" ]]; then
|
||||
echoerr "Error: you need to provide a host and port to test."
|
||||
usage
|
||||
fi
|
||||
|
||||
WAITFORIT_TIMEOUT=${WAITFORIT_TIMEOUT:-15}
|
||||
WAITFORIT_STRICT=${WAITFORIT_STRICT:-0}
|
||||
WAITFORIT_CHILD=${WAITFORIT_CHILD:-0}
|
||||
WAITFORIT_QUIET=${WAITFORIT_QUIET:-0}
|
||||
|
||||
# check to see if timeout is from busybox?
|
||||
WAITFORIT_TIMEOUT_PATH=$(type -p timeout)
|
||||
WAITFORIT_TIMEOUT_PATH=$(realpath $WAITFORIT_TIMEOUT_PATH 2>/dev/null || readlink -f $WAITFORIT_TIMEOUT_PATH)
|
||||
if [[ $WAITFORIT_TIMEOUT_PATH =~ "busybox" ]]; then
|
||||
WAITFORIT_ISBUSY=1
|
||||
WAITFORIT_BUSYTIMEFLAG="-t"
|
||||
|
||||
else
|
||||
WAITFORIT_ISBUSY=0
|
||||
WAITFORIT_BUSYTIMEFLAG=""
|
||||
fi
|
||||
|
||||
if [[ $WAITFORIT_CHILD -gt 0 ]]; then
|
||||
wait_for
|
||||
WAITFORIT_RESULT=$?
|
||||
exit $WAITFORIT_RESULT
|
||||
else
|
||||
if [[ $WAITFORIT_TIMEOUT -gt 0 ]]; then
|
||||
wait_for_wrapper
|
||||
WAITFORIT_RESULT=$?
|
||||
else
|
||||
wait_for
|
||||
WAITFORIT_RESULT=$?
|
||||
fi
|
||||
fi
|
||||
|
||||
if [[ $WAITFORIT_CLI != "" ]]; then
|
||||
if [[ $WAITFORIT_RESULT -ne 0 && $WAITFORIT_STRICT -eq 1 ]]; then
|
||||
echoerr "$WAITFORIT_cmdname: strict mode, refusing to execute subprocess"
|
||||
exit $WAITFORIT_RESULT
|
||||
fi
|
||||
exec "${WAITFORIT_CLI[@]}"
|
||||
else
|
||||
exit $WAITFORIT_RESULT
|
||||
fi
|
||||
Reference in New Issue
Block a user