From c3ef2d257ad6d80cf4c2ad511470ce63f2e41bcf Mon Sep 17 00:00:00 2001
From: Roman Masse <roman@1000mercis.com>
Date: Wed, 31 Jul 2019 15:05:21 +0200
Subject: [PATCH] Dissociate LineItem and AugmentedLineItem models.

---
 appnexus/client.py |  4 +--
 appnexus/model.py  | 62 +++++++++++++++++++++++++++++++++++++---------
 2 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/appnexus/client.py b/appnexus/client.py
index cc01f50..930dd97 100644
--- a/appnexus/client.py
+++ b/appnexus/client.py
@@ -220,8 +220,8 @@ def base_url(self):
 
 
 services_list = ["AccountRecovery", "AdProfile", "AdQualityRule", "Adserver",
-                 "Advertiser", "BatchSegment", "Brand", "Broker", "Browser",
-                 "BudgetSplitter", "Campaign", "Carrier", "Category",
+                 "Advertiser", "AugmentedLineItem", "BatchSegment", "Brand",
+                 "Broker", "Browser", "Campaign", "Carrier", "Category",
                  "ChangeLog", "ChangeLogDetail", "City", "ContentCategory",
                  "Country", "Creative", "CreativeFormat", "Currency",
                  "CustomModel", "CustomModelHash", "CustomModelLogit",
diff --git a/appnexus/model.py b/appnexus/model.py
index daaaa2b..24f922a 100644
--- a/appnexus/model.py
+++ b/appnexus/model.py
@@ -4,6 +4,7 @@
 from thingy import Thingy
 
 from appnexus.client import AppNexusClient, client, services_list
+from appnexus.exceptions import AppNexusException
 from appnexus.utils import classproperty, normalize_service_name
 
 logger = logging.getLogger("appnexus-client")
@@ -139,13 +140,6 @@ def is_ready(self):
         return (status == "ready")
 
 
-class BudgetSplitterMixin():
-
-    @property
-    def budget_splitter(self):
-        return BudgetSplitter.find_one(id=self.id)  # noqa: F821
-
-
 class ChangeLogMixin():
 
     @property
@@ -161,16 +155,60 @@ def profile(self):
         return Profile.find_one(id=self.profile_id)  # noqa: F821
 
 
+class LineItem(Model, ChangeLogMixin, ProfileMixin):
+
+    @classmethod
+    def find(cls, **kwargs):
+        kwargs["line_item_type"] = "standard_v1"
+        return super().find(**kwargs)
+
+
+class AugmentedLineItem(Model, ChangeLogMixin, ProfileMixin):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._splits = []
+
+    @classmethod
+    def find(cls, **kwargs):
+        kwargs["line_item_type"] = "standard_v2"
+        return super().find(**kwargs)
+
+    @classproperty
+    def service_name(cls):
+        return "line-item"
+
+    @property
+    def splits(self):
+        if not self._splits:
+            try:
+                self._splits = self.client.get(
+                    "budget-splitter/{}/splits".format(self.id),
+                    raw=True)
+            except AppNexusException:
+                pass
+        return self._splits
+
+    @splits.setter
+    def splits(self, value):
+        self._splits = value
+
+    def save(self, **kwargs):
+        splits = self.__dict__.pop("_splits", None)
+        super().save(**kwargs)
+        if splits:
+            self.client.modify(
+                "budget-splitter/{}/splits".format(self.id),
+                splits)
+
+
 def create_models(services_list):
     for service_name in services_list:
         ancestors = [Model]
-        if service_name in ("LineItem"):
-            ancestors.append(BudgetSplitterMixin)
-        if service_name in ("Campaign", "InsertionOrder", "LineItem",
-                            "Profile"):
+        if service_name in ("Campaign", "InsertionOrder", "Profile"):
             ancestors.append(ChangeLogMixin)
         if service_name in ("AdQualityRule", "Advertiser", "Campaign",
-                            "Creative", "LineItem", "PaymentRule"):
+                            "Creative", "PaymentRule"):
             ancestors.append(ProfileMixin)
         model = type(service_name, tuple(ancestors), {})
         globals().setdefault(service_name, model)