1"""
2Unit tests for the Inventory class.
3"""
4__copyright__ = "Copyright (C) 2014-2016  Martin Blais"
5__license__ = "GNU GPLv2"
6
7import datetime
8import unittest
9import copy
10from datetime import date
11
12from beancount.core.number import D
13from beancount.core.amount import A
14from beancount.core import amount
15from beancount.core.position import Position
16from beancount.core.position import Cost
17from beancount.core.inventory import Inventory
18from beancount.core.inventory import Booking
19from beancount.core import convert
20from beancount.core import position
21from beancount.core import inventory
22from beancount.utils import invariants
23
24
25P = position.from_string
26I = inventory.from_string
27
28
29def setUp(module):
30    invariants.instrument_invariants(Inventory,
31                                     inventory.check_invariants,
32                                     inventory.check_invariants)
33
34def tearDown(module):
35    invariants.uninstrument_invariants(Inventory)
36
37
38class TestInventory(unittest.TestCase):
39
40    def checkAmount(self, inventory, number, currency):
41        if isinstance(number, str):
42            number = D(number)
43        amount_ = amount.Amount(number, currency)
44        inv_amount = inventory.get_currency_units(amount_.currency)
45        self.assertEqual(inv_amount, amount_)
46
47    def test_from_string(self):
48        inv = inventory.from_string('')
49        self.assertEqual(Inventory(), inv)
50
51        inv = inventory.from_string('10 USD')
52        self.assertEqual(
53            Inventory([Position(A("10 USD"))]),
54            inv)
55
56        inv = inventory.from_string(' 10.00  USD ')
57        self.assertEqual(
58            Inventory([Position(A("10 USD"))]),
59            inv)
60
61        inv = inventory.from_string('1 USD, 2 CAD')
62        self.assertEqual(
63            Inventory([Position(A("1 USD")),
64                       Position(A("2 CAD"))]),
65            inv)
66
67        inv = inventory.from_string('2.2 HOOL {532.43 USD}, 3.413 EUR')
68        self.assertEqual(
69            Inventory([Position(A("2.2 HOOL"), Cost(D('532.43'), 'USD', None, None)),
70                       Position(A("3.413 EUR"))]),
71            inv)
72
73        inv = inventory.from_string(
74            '2.2 HOOL {532.43 USD}, 2.3 HOOL {564.00 USD, 2015-07-14}, 3.413 EUR')
75        self.assertEqual(
76            Inventory([Position(A("2.2 HOOL"),
77                                Cost(D('532.43'), 'USD', None, None)),
78                       Position(A("2.3 HOOL"),
79                                Cost(D('564.00'), 'USD', datetime.date(2015, 7, 14), None)),
80                       Position(A("3.413 EUR"))]),
81            inv)
82
83        inv = inventory.from_string(
84            '1.1 HOOL {500.00 # 11.00 USD}, 100 CAD')
85        self.assertEqual(
86            Inventory([Position(A("1.1 HOOL"),
87                                Cost(D('510.00'), 'USD', None, None)),
88                       Position(A("100 CAD"))]),
89            inv)
90
91    def test_ctor_empty_len(self):
92        # Test regular constructor.
93        inv = Inventory()
94        self.assertTrue(inv.is_empty())
95        self.assertEqual(0, len(inv))
96
97        inv = Inventory([P('100.00 USD'),
98                         P('101.00 USD')])
99        self.assertFalse(inv.is_empty())
100        self.assertEqual(1, len(inv))
101
102        inv = Inventory([P('100.00 USD'),
103                         P('100.00 CAD')])
104        self.assertFalse(inv.is_empty())
105        self.assertEqual(2, len(inv))
106
107        inv = Inventory()
108        self.assertEqual(0, len(inv))
109        inv.add_amount(A('100 USD'))
110        self.assertEqual(1, len(inv))
111        inv.add_amount(A('100 CAD'))
112        self.assertEqual(2, len(inv))
113
114    def test_str(self):
115        inv = I('100.00 USD, 101.00 CAD')
116        self.assertEqual('(100.00 USD, 101.00 CAD)', str(inv))
117
118    def test_copy(self):
119        inv = Inventory()
120        inv.add_amount(A('100.00 USD'))
121        self.checkAmount(inv, '100', 'USD')
122
123        # Test copying.
124        inv2 = copy.copy(inv)
125        inv2.add_amount(A('50.00 USD'))
126        self.checkAmount(inv2, '150', 'USD')
127
128        # Check that the original object is not modified.
129        self.checkAmount(inv, '100', 'USD')
130
131    def test_op_eq(self):
132        inv1 = I('100 USD, 100 CAD')
133        inv2 = I('100 CAD, 100 USD')
134        self.assertEqual(inv1, inv2)
135        self.assertEqual(inv2, inv1)
136
137        inv3 = I('200 USD, 100 CAD')
138        self.assertNotEqual(inv1, inv3)
139        self.assertNotEqual(inv3, inv1)
140
141        inv4 = I('100 USD, 100 JPY')
142        self.assertNotEqual(inv1, inv4)
143        self.assertNotEqual(inv4, inv1)
144
145        inv5 = I('100 JPY, 100 USD')
146        self.assertEqual(inv4, inv5)
147
148    def test_op_lt(self):
149        inv1 = I('100 USD, 100 CAD')
150        inv2 = I('100 CAD, 100 USD')
151        self.assertFalse(inv1 < inv2)
152        self.assertFalse(inv2 < inv1)
153
154        inv3 = I('200 USD, 100 CAD')
155        self.assertTrue(inv1 < inv3)
156        self.assertTrue(inv2 < inv3)
157        self.assertFalse(inv3 < inv1)
158
159        inv4 = I('100 USD, 100 JPY')
160        self.assertTrue(inv1 > inv4)
161
162        inv5 = I('100 JPY, 100 USD')
163        self.assertTrue(inv1 > inv5)
164        self.assertFalse(inv4 < inv5)
165
166    def test_is_small__value(self):
167        test_inv = I('1.50 JPY, 1.51 USD, 1.52 CAD')
168        for inv in test_inv, -test_inv:
169            self.assertFalse(inv.is_small(D('1.49')))
170            self.assertFalse(inv.is_small(D('1.50')))
171            self.assertTrue(inv.is_small(D('1.53')))
172            self.assertTrue(inv.is_small(D('1.52')))
173
174    def test_is_small__dict(self):
175        test_inv = I('0.03 JPY, 0.003 USD')
176        for inv in test_inv, -test_inv:
177            # Test all four types of inequalities.
178            self.assertTrue(inv.is_small({'JPY': D('0.05'), 'USD': D('0.005')}))
179            self.assertFalse(inv.is_small({'JPY': D('0.005'), 'USD': D('0.0005')}))
180            self.assertTrue(inv.is_small({'JPY': D('0.05'), 'USD': D('0.5')}))
181            self.assertFalse(inv.is_small({'JPY': D('0.005'), 'USD': D('0.005')}))
182
183            # Test border case and an epsilon under.
184            self.assertTrue(inv.is_small({'JPY': D('0.03'), 'USD': D('0.003')}))
185            self.assertFalse(inv.is_small({'JPY': D('0.02999999999999'),
186                                           'USD': D('0.003')}))
187            self.assertFalse(inv.is_small({'JPY': D('0.03'), 'USD': D('0.00299999')}))
188
189            # Test missing precisions.
190            self.assertFalse(inv.is_small({'JPY': D('0.05')}))
191            self.assertFalse(inv.is_small({'USD': D('0.005')}))
192
193            # Test extra precisions.
194            self.assertTrue(inv.is_small({'JPY': D('0.05'),
195                                          'USD': D('0.005'),
196                                          'CAD': D('0.0005')}))
197
198            # Test no precisions.
199            self.assertFalse(inv.is_small({}))
200
201    def test_is_small__with_default(self):
202        inv = I('0.03 JPY')
203        self.assertTrue(inv.is_small({'JPY': D('0.05')}))
204        self.assertFalse(inv.is_small({'JPY': D('0.02')}))
205
206    def test_is_mixed(self):
207        inv = I('100 HOOL {250 USD}, 101 HOOL {251 USD}')
208        self.assertFalse(inv.is_mixed())
209
210        inv = I('100 HOOL {250 USD}, -1 HOOL {251 USD}')
211        self.assertTrue(inv.is_mixed())
212
213        inv = I('-2 HOOL {250 USD}, -1 HOOL {251 USD}')
214        self.assertFalse(inv.is_mixed())
215
216    def test_is_reduced_by(self):
217        # Test with regular all position inventory.
218        inv = I('100 HOOL {250 USD}, 101 HOOL {251 USD}')
219        self.assertFalse(inv.is_reduced_by(A('2 HOOL')))
220        self.assertFalse(inv.is_reduced_by(A('0 HOOL')))
221        self.assertTrue(inv.is_reduced_by(A('-2 HOOL')))
222
223        # Test with a mixed-sign inventory.
224        inv = I('100 HOOL {250 USD}, -101 HOOL {251 USD}')
225        self.assertTrue(inv.is_reduced_by(A('2 HOOL')))
226        self.assertFalse(inv.is_reduced_by(A('0 HOOL')))
227        self.assertTrue(inv.is_reduced_by(A('-2 HOOL')))
228
229    def test_op_neg(self):
230        inv = Inventory()
231        inv.add_amount(A('10 USD'))
232        ninv = -inv
233        self.checkAmount(ninv, '-10', 'USD')
234
235        pinv = I('1.50 JPY, 1.51 USD, 1.52 CAD')
236        ninv = I('-1.50 JPY, -1.51 USD, -1.52 CAD')
237        self.assertEqual(pinv, -ninv)
238
239    def test_op_abs(self):
240        inv = I('1.50 USD, 2.00 USD, -1.52 CAD')
241        self.assertEqual(abs(inv), I('3.50 USD, 1.52 CAD'))
242
243    def test_op_mul(self):
244        inv = I('10 HOOL {1.11 USD}, 2.22 CAD')
245        inv2 = inv * D('3')
246        self.assertEqual(I('30 HOOL {1.11 USD}, 6.66 CAD'), inv2)
247
248    def test_get_only_position(self):
249        inv = I('10 HOOL {1.11 USD}, 2.22 CAD')
250        with self.assertRaises(AssertionError):
251            inv.get_only_position()
252        inv = I('10 HOOL {1.11 USD}')
253        self.assertEqual(A('10 HOOL'), inv.get_only_position().units)
254        inv = I('')
255        self.assertIsNone(inv.get_only_position())
256
257    def test_get_currency_units(self):
258        inv = I('40.50 JPY, 40.51 USD {1.01 CAD}, 40.52 CAD')
259        self.assertEqual(inv.get_currency_units('JPY'), A('40.50 JPY'))
260        self.assertEqual(inv.get_currency_units('USD'), A('40.51 USD'))
261        self.assertEqual(inv.get_currency_units('CAD'), A('40.52 CAD'))
262        self.assertEqual(inv.get_currency_units('AUD'), A('0 AUD'))
263        self.assertEqual(inv.get_currency_units('NZD'), A('0 NZD'))
264
265    def test_segregate_units(self):
266        inv = I('2.2 HOOL {532.43 USD}, '
267                '2.3 HOOL {564.00 USD, 2015-07-14}, '
268                '3.41 CAD, 101.20 USD')
269        ccymap = inv.segregate_units(['HOOL', 'USD', 'EUR'])
270        self.assertEqual({
271            None: I('3.41 CAD'),
272            'USD': I('101.20 USD'),
273            'EUR': inventory.Inventory(),
274            'HOOL': I('2.2 HOOL {532.43 USD}, '
275                      '2.3 HOOL {564.00 USD, 2015-07-14}')}, ccymap)
276
277    def test_units1(self):
278        inv = Inventory()
279        self.assertEqual(inv.reduce(convert.get_units), I(''))
280
281        inv = I('40.50 JPY, 40.51 USD {1.01 CAD}, 40.52 CAD')
282        self.assertEqual(inv.reduce(convert.get_units),
283                         I('40.50 JPY, 40.51 USD, 40.52 CAD'))
284
285        # Check that the same units coalesce.
286        inv = I('2 HOOL {400 USD}, 3 HOOL {410 USD}')
287        self.assertEqual(inv.reduce(convert.get_units), I('5 HOOL'))
288
289        inv = I('2 HOOL {400 USD}, -3 HOOL {410 USD}')
290        self.assertEqual(inv.reduce(convert.get_units), I('-1 HOOL'))
291
292    POSITIONS_ALL_KINDS = [
293        P('40.50 USD'),
294        P('40.50 USD {1.10 CAD}'),
295        P('40.50 USD {1.10 CAD, 2012-01-01}')]
296
297    def test_units(self):
298        inv = Inventory(self.POSITIONS_ALL_KINDS +
299                        [P('50.00 CAD')])
300        inv_cost = inv.reduce(convert.get_units)
301        self.assertEqual(I('121.50 USD, 50.00 CAD'), inv_cost)
302
303    def test_cost(self):
304        inv = Inventory(self.POSITIONS_ALL_KINDS +
305                        [P('50.00 CAD')])
306        inv_cost = inv.reduce(convert.get_cost)
307        self.assertEqual(I('40.50 USD, 139.10 CAD'), inv_cost)
308
309    def test_average(self):
310        # Identity, no aggregation.
311        inv = I('40.50 JPY, 40.51 USD {1.01 CAD}, 40.52 CAD')
312        self.assertEqual(inv.average(), inv)
313
314        # Identity, no aggregation, with a mix of lots at cost and without cost.
315        inv = I('40 USD {1.01 CAD}, 40 USD')
316        self.assertEqual(inv.average(), inv)
317
318        # Aggregation.
319        inv = I('40 USD {1.01 CAD}, 40 USD {1.02 CAD}')
320        self.assertEqual(inv.average(), I('80.00 USD {1.015 CAD}'))
321
322        # Aggregation, more units.
323        inv = I('2 HOOL {500 USD}, 3 HOOL {520 USD}, 4 HOOL {530 USD}')
324        self.assertEqual(inv.average(), I('9 HOOL {520 USD}'))
325
326        # Average on zero amount, same costs
327        inv = I('2 HOOL {500 USD}')
328        inv.add_amount(A('-2 HOOL'), Cost(D('500'), 'USD', None, None))
329        self.assertEqual(inv.average(), I(''))
330
331        # Average on zero amount, different costs
332        inv = I('2 HOOL {500 USD}')
333        inv.add_amount(A('-2 HOOL'),
334                       Cost(D('500'), 'USD', datetime.date(2000, 1, 1), None))
335        self.assertEqual(inv.average(), I(''))
336
337    def test_currencies(self):
338        inv = Inventory()
339        self.assertEqual(set(), inv.currencies())
340
341        inv = I('40 USD {1.01 CAD}, 40 USD')
342        self.assertEqual({'USD'}, inv.currencies())
343
344        inv = I('40 AAPL {1.01 USD}, 10 HOOL {2.02 USD}')
345        self.assertEqual({'AAPL', 'HOOL'}, inv.currencies())
346
347    def test_currency_pairs(self):
348        inv = Inventory()
349        self.assertEqual(set(), inv.currency_pairs())
350
351        inv = I('40 USD {1.01 CAD}, 40 USD')
352        self.assertEqual(set([('USD', 'CAD'), ('USD', None)]), inv.currency_pairs())
353
354        inv = I('40 AAPL {1.01 USD}, 10 HOOL {2.02 USD}')
355        self.assertEqual(set([('AAPL', 'USD'), ('HOOL', 'USD')]), inv.currency_pairs())
356
357    def test_add_amount(self):
358        inv = Inventory()
359        inv.add_amount(A('100.00 USD'))
360        self.checkAmount(inv, '100', 'USD')
361
362        # Add some amount
363        inv.add_amount(A('25.01 USD'))
364        self.checkAmount(inv, '125.01', 'USD')
365
366        # Subtract some amount.
367        inv.add_amount(A('-12.73 USD'))
368        self.checkAmount(inv, '112.28', 'USD')
369
370        # Subtract some to be negative (should be allowed if no lot).
371        inv.add_amount(A('-120 USD'))
372        self.checkAmount(inv, '-7.72', 'USD')
373
374        # Subtract some more.
375        inv.add_amount(A('-1 USD'))
376        self.checkAmount(inv, '-8.72', 'USD')
377
378        # Add to above zero again
379        inv.add_amount(A('18.72 USD'))
380        self.checkAmount(inv, '10', 'USD')
381
382    def test_add_amount__zero(self):
383        inv = Inventory()
384        inv.add_amount(A('0 USD'))
385        self.assertEqual(0, len(inv))
386
387    def test_add_amount__booking(self):
388        inv = Inventory()
389        _, booking = inv.add_amount(A('100.00 USD'))
390        self.assertEqual(Booking.CREATED, booking)
391
392        _, booking = inv.add_amount(A('20.00 USD'))
393        self.assertEqual(Booking.AUGMENTED, booking)
394
395        _, booking = inv.add_amount(A('-20 USD'))
396        self.assertEqual(Booking.REDUCED, booking)
397
398        _, booking = inv.add_amount(A('-100 USD'))
399        self.assertEqual(Booking.REDUCED, booking)
400
401    def test_add_amount__multi_currency(self):
402        inv = Inventory()
403        inv.add_amount(A('100 USD'))
404        inv.add_amount(A('100 CAD'))
405        self.checkAmount(inv, '100', 'USD')
406        self.checkAmount(inv, '100', 'CAD')
407
408        inv.add_amount(A('25 USD'))
409        self.checkAmount(inv, '125', 'USD')
410        self.checkAmount(inv, '100', 'CAD')
411
412    def test_add_amount__withlots(self):
413        # Testing the strict case where everything matches, with only a cost.
414        inv = Inventory()
415        inv.add_amount(A('50 HOOL'), Cost(D('700'), 'USD', None, None))
416        self.checkAmount(inv, '50', 'HOOL')
417
418        inv.add_amount(A('-40 HOOL'), Cost(D('700'), 'USD', None, None))
419        self.checkAmount(inv, '10', 'HOOL')
420
421        position_, _ = inv.add_amount(A('-12 HOOL'),
422                                      Cost(D('700'), 'USD', None, None))
423        self.assertTrue(next(iter(inv)).is_negative_at_cost())
424
425        # Testing the strict case where everything matches, a cost and a lot-date.
426        inv = Inventory()
427        inv.add_amount(A('50 HOOL'), Cost(D('700'), 'USD', date(2000, 1, 1), None))
428        self.checkAmount(inv, '50', 'HOOL')
429
430        inv.add_amount(A('-40 HOOL'), Cost(D('700'), 'USD', date(2000, 1, 1), None))
431        self.checkAmount(inv, '10', 'HOOL')
432
433        position_, _ = inv.add_amount(A('-12 HOOL'), Cost(D('700'), 'USD',
434                                                          date(2000, 1, 1), None))
435        self.assertTrue(next(iter(inv)).is_negative_at_cost())
436
437    def test_add_amount__allow_negative(self):
438        inv = Inventory()
439
440        # Test adding positions of different types.
441        position_, _ = inv.add_amount(A('-11 USD'))
442        self.assertIsNone(position_)
443        position_, _ = inv.add_amount(A('-11 USD'),
444                                      Cost(D('1.10'), 'CAD', None, None))
445        self.assertIsNone(position_)
446        position_, _ = inv.add_amount(A('-11 USD'),
447                                      Cost(D('1.10'), 'CAD', date(2012, 1, 1), None))
448        self.assertIsNone(position_)
449
450        # Check for reductions.
451        invlist = list(inv)
452        self.assertTrue(invlist[1].is_negative_at_cost())
453        self.assertTrue(invlist[2].is_negative_at_cost())
454        inv.add_amount(A('-11 USD'), Cost(D('1.10'), 'CAD', None, None))
455        inv.add_amount(A('-11 USD'), Cost(D('1.10'), 'CAD', date(2012, 1, 1), None))
456        self.assertEqual(3, len(inv))
457
458        # Test adding to a position that does exist.
459        inv = I('10 USD, 10 USD {1.10 CAD}, 10 USD {1.10 CAD, 2012-01-01}')
460        position_, _ = inv.add_amount(A('-11 USD'))
461        self.assertEqual(position_, position.from_string('10 USD'))
462        position_, _ = inv.add_amount(A('-11 USD'),
463                                      Cost(D('1.10'), 'CAD', None, None))
464        self.assertEqual(position_, position.from_string('10 USD {1.10 CAD}'))
465        position_, _ = inv.add_amount(A('-11 USD'),
466                                      Cost(D('1.10'), 'CAD', date(2012, 1, 1), None))
467        self.assertEqual(position_, position.from_string('10 USD {1.10 CAD, 2012-01-01}'))
468
469    def test_add_position(self):
470        inv = Inventory()
471        for pos in self.POSITIONS_ALL_KINDS:
472            inv.add_position(pos)
473        self.assertEqual(Inventory(self.POSITIONS_ALL_KINDS), inv)
474
475    def test_op_add(self):
476        inv1 = I('17.00 USD')
477        orig_inv1 = I('17.00 USD')
478        inv2 = I('21.00 CAD')
479        inv3 = inv1 + inv2
480        self.assertEqual(I('17.00 USD, 21.00 CAD'), inv3)
481        self.assertEqual(orig_inv1, inv1)
482
483    def test_update(self):
484        inv1 = I('11 USD')
485        inv2 = I('12 CAD')
486        inv_updated = inv1.add_inventory(inv2)
487        expect_updated = I('11 USD, 12 CAD')
488        self.assertEqual(expect_updated, inv_updated)
489        self.assertEqual(expect_updated, inv1)
490
491    def test_sum_inventories(self):
492        inv1 = Inventory()
493        inv1.add_amount(A('10 USD'))
494
495        inv2 = Inventory()
496        inv2.add_amount(A('20 CAD'))
497        inv2.add_amount(A('55 HOOL'))
498
499        _ = inv1 + inv2
500
501    def test_reduce(self):
502        inv = I('100.00 USD, 101.00 CAD, 100 HOOL {300.00 USD}')
503        inv_units = inv.reduce(lambda posting: posting.units)
504        self.assertEqual(I('100.00 USD, 101.00 CAD, 100 HOOL'), inv_units)
505
506
507if __name__ == '__main__':
508    unittest.main()
509