Coverage for product_risk_suite/threat_model/migrations/0004_threatmodelconnectionname_primary_key.py: 59%

1 

2from django.db import migrations, models 

3 

4tm_cn_map = {} 

5def migrate_pop_tm_cn(apps, schema_editor): 

6 ThreatModel = apps.get_model('threat_model', 'ThreatModel') 

7 

8 for tm in ThreatModel.objects.all(): 

9 tm_cn_map[tm.id] = [] 

10 for svg_id in tm.connection_names.all(): 

11 tm_cn_map[tm.id].append(svg_id.id) 

12 tm.connection_names.clear() 

13 tm.save() 

14 

15def migrate_push_tm_cn(apps, schema_editor): 

16 ThreatModel = apps.get_model('threat_model', 'ThreatModel') 

17 ThreatModelConnectionName = apps.get_model('threat_model', 'ThreatModelConnectionName') 

18 

19 for tm in ThreatModel.objects.all(): 

20 for svg_id in tm_cn_map[tm.id]: 

21 tmcn = ThreatModelConnectionName.objects.get(id = svg_id) 

22 tm.connection_names.add(tmcn) 

23 tm.save() 

24 

25class Migration(migrations.Migration): 

26 atomic = False 

27 

28 dependencies = [ 

29 ('threat_model', '0003_threatmodelconnectionname_id'), 

30 ('product', '0005_remove_productriskentry_svg_id'), 

31 ] 

32 

33 operations = [ 

34 migrations.RunPython(migrate_pop_tm_cn), 

35 migrations.RemoveField( 

36 model_name='threatmodel', 

37 name='connection_names', 

38 ), 

39 migrations.AlterField( 

40 model_name='threatmodelconnectionname', 

41 name='id', 

42 field=models.AutoField(primary_key=True, serialize=False), 

43 ), 

44 migrations.AlterField( 

45 model_name='threatmodelconnectionname', 

46 name='name', 

47 field=models.CharField(max_length=255), 

48 ), 

49 migrations.AddField( 

50 model_name='threatmodel', 

51 name='connection_names', 

52 field=models.ManyToManyField(blank=True, to='threat_model.threatmodelconnectionname'), 

53 ), 

54 migrations.RunPython(migrate_push_tm_cn), 

55 ]