🔑 A simple database server, for getting and setting from a key-value store.

dbserver.py

#!/usr/bin/env python3

import unittest
import json
import io
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.error import HTTPError
from urllib.request import urlopen


class Handler(BaseHTTPRequestHandler):
    '''
    Handle requests!
    '''
    # Stash submitted data in a hash map in memory
    datastore = {}

    def do_GET(self):
        '''
        The handler should respond to two patterns of GET requests:

            - /get?key=keyname -- return the value stored in the key `keyname`
            - /set?keyname=value -- store the value `value` under the key `keyname`
        '''
        # Get the base command ('get' or 'set') and the params
        cmd, param = self.path.split('?')

        # Strip off the leading backslash
        cmd = cmd[1:]

        # Set some default response values
        msg = 'Bad request'
        status_code = 400
        response = None

        if cmd == 'set':
            # Set a key/value pair
            key, val = param.split('=')
            self.datastore[key] = val

            status_code = 201
            msg = 'SET request receieved for params %s=%s' % (key, val)

        elif cmd == 'get':
            # Retreive a key/value pair
            prefix, key = param.split('=')

            if prefix == 'key':

                if self.datastore.get(key):
                    status_code = 200
                    val = self.datastore[key]
                    output = {key: val}
                    response = json.dumps(output).encode('utf-8')
                else:
                    status_code = 404
                    msg = 'Key "%s" not found' % key

        self.send_response(status_code, message=msg)
        self.send_header("Content-type", "application/json")
        self.end_headers()

        if response:
            self.wfile.write(response)


class TestServer(unittest.TestCase):
    '''
    Unit tests for the request handler.

    To run tests, start by running the server:

    `python3 dbserver.py`

    Then open another shell and run the tests:

    `python3 -m unittest dbserver.py`
    '''
    def test_set_key(self):
        set_key = urlopen('http://127.0.0.1:4000/set?testkey=testval')
        self.assertEqual(set_key.status, 201)
        self.assertEqual(set_key.msg, 'SET request receieved for params testkey=testval')

    def test_set_and_get_key(self):
        set_key = urlopen('http://127.0.0.1:4000/set?testkey=testval')
        get_key = urlopen('http://127.0.0.1:4000/get?key=testkey')

        self.assertEqual(get_key.status, 200)
        self.assertEqual(get_key.read().decode('utf-8'), '{"testkey": "testval"}')

    def test_key_doesnt_exist(self):
        with self.assertRaises(HTTPError):
            get_key = urlopen('http://127.0.0.1:4000/get?key=badkey')
            self.assertEqual(get_key.msg, 'Key "badkey" not found')


if __name__ == '__main__':
    # Run a simple server on the loopback interface
    port = 4000
    print('Serving HTTP on 127.0.0.1:%d...' % port)
    server = HTTPServer(('', port), Handler)
    server.serve_forever()