diff --git a/demand_generator/tek_data/gen_heat_profile.py b/demand_generator/tek_data/gen_heat_profile.py
index cb8ced976abbcdf2c0b28ed38d28075ab349e9e2..37c363aac4fc3f8d4c4e2caf81d4fefbcf5b79e5 100644
--- a/demand_generator/tek_data/gen_heat_profile.py
+++ b/demand_generator/tek_data/gen_heat_profile.py
@@ -8,7 +8,6 @@ import datetime
 from warnings import warn
 import pandas as pd
 import numpy as np
-import matplotlib.pyplot as plt
 
 # ==============================================================================
 #                     Non-residential Building List
@@ -45,12 +44,11 @@ output_path = os.path.join(base_path, "data", "tek_data", "output_heat_profile")
 
 
 def gen_heat_profile(building_typ,
-                     area,
                      temperature_profile,
+                     yearly_demand=None,
+                     area=None,
                      year=2021,
-                     energy_typ="mittel",
-                     plot=False,
-                     save_plot=False):
+                     energy_typ="mittel"):
     """
     total_degree_day: K*h
     annual_value: kW*h, jährlicher Gesamt Heizwärmebedarf
@@ -58,6 +56,18 @@ def gen_heat_profile(building_typ,
     on room type and heating start at the temperature of 15 degree.
     :return:
     """
+    if yearly_demand is None and area is not None:
+        pass
+    elif yearly_demand is not None and area is None:
+        area = 300
+    elif yearly_demand is None and area is None:
+        warn("By heat demand generating is no yearly demand nor area for the "
+             "building given")
+        area = 300
+    elif yearly_demand is not None and area is not None:
+        warn("By heat demand generating both yearly demand and area for the "
+             "building are given. Only the yearly demand would be used.")
+
     # Analysis thermal zones in building
     new_zone_df = analysis_bld_zone(building_typ, area)
 
@@ -77,19 +87,35 @@ def gen_heat_profile(building_typ,
                                     axis=0)
         total_heat_demand += zone_heat_demand
 
-    if plot:
-        plot_profile(total_heat_profile, save_plot)
+    if yearly_demand is not None:
+        total_heat_profile = np.array(total_heat_profile) * yearly_demand / \
+                             total_heat_demand
 
     return total_heat_profile
 
 
-def gen_hot_water_profile(building_typ, area, year=2021, energy_typ="mittel"):
+def gen_hot_water_profile(building_typ,
+                          yearly_demand=None,
+                          area=None, year=2021, energy_typ="mittel"):
     # new_zone_df = analysis_bld_zone(building_typ, area)
     # for row in range(len(new_zone_df)):
     #     zone = new_zone_df.loc[row, 'DIN_Zone']
+    if yearly_demand is None and area is not None:
+        pass
+    elif yearly_demand is not None and area is None:
+        area = 300
+    elif yearly_demand is None and area is None:
+        warn("By hot water demand generating is no yearly demand nor area for "
+             "the building given")
+        area = 300
+    elif yearly_demand is not None and area is not None:
+        warn("By hot water demand generating both yearly demand and area for "
+             "the building are given. Only the yearly demand would be used.")
 
     bld_hot_water_demand = calc_bld_demand(building_typ, area, 'hot_water',
                                            energy_typ)
+    hot_water_heating_demand_array = None
+
     if building_typ == "Wohngebäude" or building_typ == "Wohngebäude (MFH)":
         hot_water_heating_demand_df = pd.read_excel(input_dhw_path,
                                                     sheet_name='DHW',
@@ -118,6 +144,11 @@ def gen_hot_water_profile(building_typ, area, year=2021, energy_typ="mittel"):
             bld_occupancy) * np.array(bld_occupancy)
         # print(hot_water_heating_demand_array)
 
+    if yearly_demand is not None:
+        hot_water_heating_demand_array = hot_water_heating_demand_array * \
+                                         yearly_demand / \
+                                         hot_water_heating_demand_array.sum()
+
     return hot_water_heating_demand_array
 
 
@@ -317,8 +348,9 @@ def calc_bld_occupancy(building_typ,
 
             for day in range(365):
                 for hour in range(24):
-                    if hour_status[day*24+hour] == 1:
-                        zone_occupancy[day*24+hour] = zone_pers_profile[hour]
+                    if hour_status[day * 24 + hour] == 1:
+                        zone_occupancy[day * 24 + hour] = zone_pers_profile[
+                            hour]
             # print(zone_pers_profile)
             bld_occupancy += np.array(zone_occupancy)
     # print(bld_occupancy)
@@ -354,19 +386,6 @@ def degree_day(zone_typ, annual_value, profile_df, temperature_profile,
     return heat_profile
 
 
-def plot_profile(heat_profile, save_plot=False):
-    plt.figure()
-    plt.plot(heat_profile)
-    plt.ylabel('Heat Profile')
-    plt.xlabel('Hours [h]')
-    plt.ylim(ymin=0)
-    plt.xlim(xmin=0)
-    plt.grid()
-    if save_plot:
-        plt.savefig(os.path.join(output_path, 'heat_profile_figure.jpg'))
-    plt.show()
-
-
 def calc_residential_demand(bld_type, bld_year, bld_area,
                             method='TABULA Berechnungsverfahren / korrigiert '
                                    'auf Niveau von Verbrauchswerten',
@@ -391,13 +410,88 @@ def calc_residential_demand(bld_type, bld_year, bld_area,
     return heating_demand, hot_water_demand
 
 
-if __name__ == "__main__":
+def test_gen_heat_profile():
     input_temp_path = os.path.join(base_path, 'tek_data', 'temperature.csv')
     temperature = pd.read_csv(input_temp_path)['temperature'].values
+
+    assert isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature,
+                                       area=250), list) or \
+           isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature,
+                                       area=250), np.ndarray), \
+        "The return value type is " + str(type(gen_heat_profile(
+            "Wohngebäude (MFH)", temperature, area=250)))
+    assert isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature,
+                                       yearly_demand=50000), list) or \
+           isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature,
+                                       yearly_demand=50000), np.ndarray), \
+        "The return value type is " + str(type(gen_heat_profile(
+            "Wohngebäude (MFH)", temperature, yearly_demand=50000)))
+
+    assert isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature,
+                                       yearly_demand=50000, area=250), list) or \
+           isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature,
+                                       yearly_demand=50000, area=250),
+                      np.ndarray), \
+        "The return value type is " + str(type(gen_heat_profile(
+            "Wohngebäude (MFH)", temperature, yearly_demand=50000,
+            area=250)))
+    assert isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature),
+                      list) or \
+           isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature),
+                      np.ndarray), \
+        "The return value type is " + str(type(gen_heat_profile(
+            "Wohngebäude (MFH)", temperature)))
+
+
+def test_gen_hot_water_profile():
+    assert isinstance(gen_hot_water_profile("Wohngebäude (MFH)",
+                                            area=250), list) or \
+           isinstance(gen_hot_water_profile("Wohngebäude (MFH)",
+                                            area=250), np.ndarray), \
+           "The return value type is " + str(type(gen_hot_water_profile(
+                                                 "Wohngebäude (MFH)",
+                                             area=250)))
+    assert isinstance(gen_hot_water_profile("Wohngebäude (MFH)",
+                                            yearly_demand=50000), list) or \
+           isinstance(gen_hot_water_profile("Wohngebäude (MFH)",
+                                            yearly_demand=50000), np.ndarray), \
+           "The return value type is " + str(type(gen_hot_water_profile(
+               "Wohngebäude (MFH)", yearly_demand=50000)))
+
+    assert isinstance(gen_hot_water_profile("Wohngebäude (MFH)",
+                                            yearly_demand=50000, area=250),
+                      list) or \
+           isinstance(gen_hot_water_profile("Wohngebäude (MFH)",
+                                            yearly_demand=50000, area=250),
+                      np.ndarray), \
+           "The return value type is " + str(type(gen_hot_water_profile(
+               "Wohngebäude (MFH)", yearly_demand=50000, area=250)))
+    assert isinstance(gen_hot_water_profile("Wohngebäude (MFH)"),
+                      list) or \
+           isinstance(gen_hot_water_profile("Wohngebäude (MFH)"),
+                      np.ndarray), \
+           "The return value type is " + str(type(gen_hot_water_profile(
+               "Wohngebäude (MFH)")))
+
+
+if __name__ == "__main__":
+    # input_temp_path = os.path.join(base_path, 'tek_data', 'temperature.csv')
+    # temperature = pd.read_csv(input_temp_path)['temperature'].values
+    pass
     # gen_heat_profile("Wohngebäude", 300, temperature, plot=True)
     # print('Space heating demand profile:')
     # print(gen_heat_profile("Wohngebäude (MFH)", 300, temperature))
-    print('Hot water demand profile:')
-    print(gen_hot_water_profile("Verwaltungsgebäude", 300))
+    # print('Hot water demand profile:')
+    # print(gen_hot_water_profile("Verwaltungsgebäude", 300))
 
     # calc_bld_occupancy("Verwaltungsgebäude", 300, year=2021)
+
+    # def test_gen_heat_profile():
+    #     assert isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature,
+    #                                        area=250), list)
+    #     assert isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature,
+    #                                        yearly_demand=50000), list)
+    #     assert isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature,
+    #                                        yearly_demand=50000, area=250), list)
+    #     assert isinstance(gen_heat_profile("Wohngebäude (MFH)", temperature),
+    #                       list)